FoodDesert commited on
Commit
52ae321
·
verified ·
1 Parent(s): dc69cb1

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +744 -745
app.py CHANGED
@@ -1,745 +1,744 @@
1
- import gradio as gr
2
- from sklearn.metrics.pairwise import cosine_similarity
3
- from scipy.sparse import csr_matrix
4
- import numpy as np
5
- import joblib
6
- from joblib import load
7
- import h5py
8
- from io import BytesIO
9
- import csv
10
- import re
11
- import random
12
- import compress_fasttext
13
- from collections import OrderedDict
14
- from lark import Lark, Tree, Token
15
- from lark.exceptions import ParseError
16
- import json
17
- import zipfile
18
- from PIL import Image
19
- import io
20
- import os
21
- import glob
22
- import itertools
23
- from itertools import islice
24
- from pathlib import Path
25
- import logging
26
-
27
- # Set up logging
28
- logging.basicConfig(filename='error.log', level=logging.DEBUG, format='%(asctime)s %(levelname)s:%(message)s')
29
-
30
-
31
- faq_content="""
32
- # Questions:
33
-
34
- ## What is the purpose of this tool?
35
-
36
- Since Stable Diffusion's initial release in 2022, users have developed a myriad of fine-tuned text to image models, each with unique "linguistic" preferences depending on the data from which it was fine-tuned.
37
- Some models react best when prompted with verbose scene descriptions akin to DALL-E, while others fine-tuned on images scraped from popular image boards understand those boards' tag sets.
38
- This tool serves as a linguistic bridge to the e621 image board tag lexicon, on which many popular models such as Fluffyrock, Fluffusion, and Pony Diffusion v6 were trained.
39
-
40
- When you enter a txt2img prompt and press the "submit" button, Prompt Squirrel parses your prompt and checks that all your tags are valid e621 tags.
41
- If it finds any that are not, it recommends some valid e621 tags you can use to replace them in the "Unknown Tags" section.
42
- Additionally, in the "Top Artists" text box, it lists the artists who would most likely draw an image having the set of tags you provided.
43
- This is useful to align your prompt with the expected input to an e621-trained model.
44
-
45
- ## Does input order matter?
46
-
47
- No
48
-
49
- ## Should I use underscores or spaces in the input tags?
50
-
51
- As a rule, e621-trained models replace underscores in tags with spaces, so spaces are preferred.
52
-
53
- ## Can I use parentheses or weights as in the Stable Diffusion Automatic1111 WebUI?
54
-
55
- Yes, but only '(' and ')' and numerical weights, and all of these things are ignored in all calculations. The main benefit of this is that you can copy/paste prompts from one program to another with minimal editing.
56
- An example that illustrates acceptable parentheses and weight formatting is:
57
- ((sunset over the mountains)), (clear sky:1.5), ((eagle flying high:2.0)), river, (fish swimming in the river:1.2), (campfire, (marshmallows:2.1):1.3), stars in the sky, ((full moon:1.8)), (wolf howling:1.7)
58
-
59
- ## Why are some valid tags marked as "unknown", and why don't some artists ever get returned?
60
-
61
- Some data is excluded from consideration if it did not occur frequently enough in the sample from which the application makes its calculations.
62
- If an artist or tag is too infrequent, we might not think we have enough data to make predictions about it. Additionally, Prompt Squirrel gathers information from several sources, and the sources are not always consistent about things like exact tag names or counts, which vary over time.
63
-
64
- ## Why do some suggested tags not have summaries or wiki links, and of those that do, why do some look truncated?
65
-
66
- Both of these features are extracted from the tag wiki pages, but some valid e621 tags do not have wiki pages. Additionally, the summaries are heuristically extracted from the beginning of the wiki pages, and this extraction process is prone to some amount of error.
67
-
68
- ## Are there any special tags?
69
-
70
- Yes. We normalized the favorite counts of each image to a range of 0-9, with 0 being the lowest favcount, and 9 being the highest.
71
- You can include any of these special tags: "score:0", "score:1", "score:2", "score:3", "score:4", "score:5", "score:6", "score:7", "score:8", "score:9"
72
- in your list to bias the output toward artists with higher or lower scoring images.
73
-
74
- ## Are there any other special tricks?
75
-
76
- Yes. If you want to more strongly bias the artist output toward a specific tag, you can just list it multiple times.
77
- So for example, the query "red fox, red fox, red fox, score:7" will yield a list of artists who are more strongly associated with the tag "red fox"
78
- than the query "red fox, score:7".
79
-
80
- ## Why is this space tagged "not-for-all-audience"
81
- The "not-for-all-audience" tag informs users that this tool's text output is derived from e621.net data for tag prediction and completion.
82
- The app will try not to display nsfw tags unless the "Allow NSFW Tags" is checked, but the filter is not perfect.
83
-
84
- ## How is the artist list calculated?
85
-
86
- Each artist is represented by a "pseudo-document" composed of all the tags from their uploaded images, treating these tags similarly to words in a text document.
87
- Similarly, when you input a set of tags, the system creates a pseudo-document for your query out of all the tags.
88
- It then compares your tags against each artist's collection, essentially finding which artist's tags are most "similar" to yours.
89
- This method helps identify artists whose work is closely aligned with the themes or elements you're interested in.
90
- For those curious about the underlying mechanics of comparing text-like data, we employ the TF-IDF (Term Frequency-Inverse Document Frequency) method, a standard approach in information retrieval, and reduce the TF-IDF matrix to a reasonable size using Singular Value Decomposition.
91
- You can read more about TF-IDF on its [Wikipedia page](https://en.wikipedia.org/wiki/Tf%E2%80%93idf) and Singular Value Decomposition on its [Wikipedia page](https://en.wikipedia.org/wiki/Singular_value_decomposition).
92
-
93
- ## How does the tag corrector work?
94
-
95
- We collect the tag sets from over 4 million e621 posts, treating the tag set from each image as an individual document.
96
- We then randomly replace about 10% of the tags in each document with a randomly selected alias from e621's list of aliases for the tag
97
- (e.g. "canine" gets replaced with one of {k9,canines,mongrel,cannine,cnaine,feral_canine,anthro_canine}).
98
- We then train a FastText (https://fasttext.cc/) model on the documents. The result of this training is a function that maps arbitrary words to vectors such that
99
- the vector for a tag and the vectors for its aliases are all close together (because the model has seen them in similar contexts).
100
- Since the lists of aliases contain misspellings and rephrasings of tags, the model should be robust to these kinds of problems as long as they are not too dissimilar from the alias lists.
101
-
102
- To enhance the tag corrector further, we employ the same TF-IDF method we used for artist tags to calculate a separate, context-sensitive similarity score for each of the top 100 tags selected via the FastText method.
103
- By considering the context in which tags are used, we can now not only correct misspellings and rephrasings but also make more contextually relevant suggestions.
104
- The "similarity weight" slider controls how much weight these TF-IDF scores are given vs how much weight the FastText similarity model is given when suggesting replacements for invalid tags.
105
- A similarity weight slider value of 0 means that only the FastText model's predictions will be used to calculate similarity scores, and a value of 1 means only the TF-IDF scores are used (although the FastText model is still used to trim the list of candidates).
106
-
107
-
108
- ## How do the sample images work?
109
-
110
- In the first row of galleries, for each artist in the dataset, we generated a sample image with the model Fluffyrock Unleashed using the prompt "by artist, soyjak, anthro, male, bust portrait, meme, grin" where "artist" is the name of an artist.
111
- The simplicity of the prompt, the the simplicty of the default style, and the recognizability of the character make it easier to understand how artist names affect generated image styles.
112
- The image on the left captioned "No Artist" was generated with the same prompt, but with no artist name.
113
- You should compare all the images to the first to see how the artist names affect the output.
114
- Each subsequent row of images was generated using the same process, but with a different prompt.
115
- See SamplePrompts.csv for the list of prompts used and their descriptions.
116
- """
117
-
118
-
119
- nsfw_threshold = 0.95 # Assuming the threshold value is defined here
120
-
121
- css = """
122
- .scrollable-content {
123
- max-height: 500px;
124
- overflow-y: auto;
125
- }
126
- """
127
-
128
- grammar=r"""
129
- !start: (prompt | /[][():]/+)*
130
- prompt: (emphasized | plain | comma | WHITESPACE)*
131
- !emphasized: "(" prompt ")"
132
- | "(" prompt ":" [WHITESPACE] NUMBER [WHITESPACE] ")"
133
- comma: ","
134
- WHITESPACE: /\s+/
135
- plain: /([^,\\\[\]():|]|\\.)+/
136
- %import common.SIGNED_NUMBER -> NUMBER
137
- """
138
-
139
- # Initialize the parser
140
- parser = Lark(grammar, start='start')
141
-
142
- # Function to extract tags
143
- def extract_tags(tree):
144
- tags_with_positions = []
145
- def _traverse(node):
146
- if isinstance(node, Token) and node.type == '__ANON_1':
147
- tag_position = node.start_pos
148
- tag_text = node.value
149
- tags_with_positions.append((tag_text, tag_position, "tag"))
150
- elif not isinstance(node, Token):
151
- for child in node.children:
152
- _traverse(child)
153
- _traverse(tree)
154
- return tags_with_positions
155
-
156
-
157
- special_tags = ["score:0", "score:1", "score:2", "score:3", "score:4", "score:5", "score:6", "score:7", "score:8", "score:9", "rating:s", "rating:q", "rating:e"]
158
- def remove_special_tags(original_string):
159
- tags = [tag.strip() for tag in original_string.split(",")]
160
- remaining_tags = [tag for tag in tags if tag not in special_tags]
161
- removed_tags = [tag for tag in tags if tag in special_tags]
162
- return ", ".join(remaining_tags), removed_tags
163
-
164
-
165
- # Define a function to load all necessary components
166
- def load_model_components(file_path):
167
- # Ensure the file path is a Path object for robust path handling
168
- file_path = Path(file_path)
169
-
170
- # Check if the file exists
171
- if not file_path.is_file():
172
- raise FileNotFoundError(f"The specified joblib file was not found: {file_path}")
173
-
174
- # Load all the model components from the joblib file
175
- model_components = joblib.load(file_path)
176
-
177
- # Create a reverse mapping from row index to tag
178
- if 'tag_to_row_index' in model_components:
179
- model_components['row_to_tag'] = {idx: tag for tag, idx in model_components['tag_to_row_index'].items()}
180
-
181
- return model_components
182
-
183
- # Load all components at the start
184
- tf_idf_components = load_model_components('tf_idf_files_420.joblib')
185
-
186
-
187
- nsfw_tags = set() # Initialize an empty set to store words meeting the threshold
188
- # Open and read the CSV file
189
- with open("word_rating_probabilities.csv", 'r', newline='', encoding='utf-8') as csvfile:
190
- reader = csv.reader(csvfile)
191
- next(reader, None) # Skip the header row
192
- for row in reader:
193
- word = row[0] # The word is in the first column
194
- probability_sum = float(row[1]) # The sum of probabilities is in the second column, convert to float for comparison
195
- # Check if the probability sum meets the threshold and add the word to the set if it does
196
- if probability_sum >= nsfw_threshold:
197
- nsfw_tags.add(word)
198
-
199
-
200
- # Read the set of valid artists into memory.
201
- artist_set = set()
202
- with open("fluffyrock_3m.csv", 'r', newline='', encoding='utf-8') as csvfile:
203
- """
204
- Load artist names from a CSV file and store them in the global set.
205
- Artist tags start with 'by_' and the prefix will be removed.
206
- """
207
- reader = csv.reader(csvfile)
208
- for row in reader:
209
- tag_name = row[0] # Assuming the first column contains the tag names
210
- if tag_name.startswith('by_'):
211
- # Strip 'by_' from the start of the tag name and add to the set
212
- artist_name = tag_name[3:] # Remove the first three characters 'by_'
213
- artist_set.add(artist_name)
214
- def is_artist(name):
215
- return name in artist_set
216
-
217
-
218
- sample_images_directory_path = 'sampleimages'
219
- def generate_artist_image_tuples(top_artists, image_directory):
220
- json_files = glob.glob(f'{image_directory}/*.json')
221
- json_file_path = json_files[0] if json_files else None
222
- with open(json_file_path, 'r') as json_file:
223
- artist_to_file_map = json.load(json_file)
224
-
225
- filename = artist_to_file_map.get("")
226
- image_path = os.path.join(image_directory, filename)
227
- if os.path.exists(image_path):
228
- baseline_tuple = [(image_path, "No Artist")]
229
-
230
- artist_image_tuples = []
231
- for artist in top_artists:
232
- filename = artist_to_file_map.get(artist)
233
- if filename:
234
- image_path = os.path.join(image_directory, filename)
235
- if os.path.exists(image_path):
236
- artist_image_tuples.append((image_path, artist if artist else "No Artist"))
237
-
238
- return baseline_tuple, artist_image_tuples
239
-
240
-
241
- def clean_tag(tag):
242
- return ''.join(char for char in tag if ord(char) < 128)
243
-
244
-
245
- #Normally returns tag to aliases, but when reverse=True, returns alias to tags
246
- def build_aliases_dict(filename, reverse=False):
247
- aliases_dict = {}
248
- with open(filename, 'r', newline='', encoding='utf-8') as csvfile:
249
- reader = csv.reader(csvfile)
250
- for row in reader:
251
- tag = clean_tag(row[0])
252
- alias_list = [] if row[3] == "null" else [clean_tag(alias) for alias in row[3].split(',')]
253
- if reverse:
254
- for alias in alias_list:
255
- aliases_dict.setdefault(alias, []).append(tag)
256
- else:
257
- aliases_dict[tag] = alias_list
258
- return aliases_dict
259
-
260
-
261
- def build_tag_count_dict(filename):
262
- with open(filename, 'r', newline='', encoding='utf-8') as csvfile:
263
- reader = csv.reader(csvfile)
264
- result_dict = {}
265
- for row in reader:
266
- key = row[0]
267
- value = int(row[2]) if row[2].isdigit() else None
268
- if value is not None:
269
- result_dict[key] = value
270
- return result_dict
271
-
272
- import csv
273
-
274
-
275
- def build_tag_id_wiki_dict(filename='wiki_pages-2023-08-08.csv'):
276
- """
277
- Reads a CSV file and returns a dictionary mapping tag names to tuples of
278
- (number, most relevant line from the wiki entry). Rows with a non-integer in the first column are ignored.
279
- The most relevant line is the first line that does not start with "thumb" and is not blank.
280
-
281
- Parameters:
282
- - filename: The path to the CSV file.
283
-
284
- Returns:
285
- - A dictionary where each key is a tag name and each value is a tuple (number, most relevant wiki entry line).
286
- """
287
- tag_data = {}
288
- with open(filename, 'r', encoding='utf-8') as csvfile:
289
- reader = csv.reader(csvfile)
290
-
291
- # Skip the header row
292
- next(reader)
293
-
294
- for row in reader:
295
- try:
296
- # Attempt to convert the first column to an integer
297
- number = int(row[0])
298
- except ValueError:
299
- # If conversion fails, skip this row
300
- continue
301
-
302
- tag = row[3]
303
- wiki_entry_full = row[4]
304
-
305
- # Process the wiki_entry to find the most relevant line
306
- relevant_line = ''
307
- for line in wiki_entry_full.split('\n'):
308
- if line.strip() and not line.startswith("thumb"):
309
- relevant_line = line
310
- break
311
-
312
- # Map the tag to a tuple of (number, relevant_line)
313
- tag_data[tag] = (number, relevant_line)
314
-
315
- return tag_data
316
-
317
-
318
- def create_html_tables_for_tags(subtable_heading, item_heading, word_similarity_tuples, tag2count, tag2idwiki):
319
- # Wrap the tag part in a <span> with styles for bold and larger font
320
- html_str = f"<div style='display: inline-block; margin: 20px; vertical-align: top;'><table><thead><tr><th colspan='3' style='text-align: center; padding-bottom: 10px;'><span style='font-weight: bold; font-size: 20px;'>{subtable_heading}</span></th></tr></thead><tbody><tr style='border-bottom: 1px solid #000;'><th>{item_heading}</th><th>Similarity</th><th>Count</th></tr>"
321
- # Loop through the results and add table rows for each
322
- for word, sim in word_similarity_tuples:
323
- word_with_underscores = word.replace(' ', '_')
324
- word_with_escaped_parentheses = word.replace("\\(", "(").replace("\\)", ")").replace("(", "\\(").replace(")", "\\)")
325
- count = tag2count.get(word_with_underscores.replace("\\(", "(").replace("\\)", ")"), 0) # Get the count if available, otherwise default to 0
326
- tag_id, wiki_entry = tag2idwiki.get(word_with_underscores, (None, ''))
327
- # Check if tag_id and wiki_entry are valid
328
- if tag_id is not None and wiki_entry:
329
- # Construct the URL for the tag's wiki page
330
- wiki_url = f"https://e621.net/wiki_pages/{tag_id}"
331
- # Make the tag a hyperlink with a tooltip
332
- tag_element = f"<a href='{wiki_url}' target='_blank' title='{wiki_entry}'>{word_with_escaped_parentheses}</a>"
333
- else:
334
- # Display the word without any hyperlink or tooltip
335
- tag_element = word_with_escaped_parentheses
336
- # Include the tag element in the table row
337
- html_str += f"<tr><td style='border: none; padding: 5px; height: 20px;'>{tag_element}</td><td style='border: none; padding: 5px; height: 20px;'>{round(sim, 3)}</td><td style='border: none; padding: 5px; height: 20px;'>{count}</td></tr>"
338
-
339
- html_str += "</tbody></table></div>"
340
- return html_str
341
-
342
-
343
- def create_top_artists_table(top_artists):
344
- # Add a heading above the table
345
- html_str = "<div class=\"scrollable-content\" style='display: inline-block; margin: 20px; text-align: center;'>"
346
- html_str += "<h1>Top Artists</h1>" # Heading for the table
347
- # Start the table with increased font size and no borders between rows
348
- html_str += "<table style='font-size: 20px; border-collapse: collapse;'>"
349
- html_str += "<thead><tr><th>Artist</th><th>Similarity</th></tr></thead><tbody>"
350
- # Loop through the top artists and add a row for each without the rank and without borders between rows
351
- for artist, score in top_artists:
352
- artist_name = artist[3:] if artist.startswith("by ") else artist # Remove "by " prefix
353
- similarity_percentage = "{:.1f}%".format(score * 100) # Convert score to percentage string with one decimal
354
- html_str += f"<td style='padding: 3px 20px; border: none;'>{artist_name}</td><td style='padding: 3px 20px; border: none;'>{similarity_percentage}</td></tr>"
355
-
356
- # Close the table HTML
357
- html_str += "</tbody></table></div>"
358
-
359
- return html_str
360
-
361
-
362
- def construct_pseudo_vector(pseudo_doc_terms, idf_loaded, tag_to_row_loaded):
363
- # Initialize a vector of zeros with the length of the term_to_index mapping
364
- pseudo_vector = np.zeros(len(tag_to_row_loaded))
365
-
366
- # Fill in the vector for terms in the pseudo document
367
- for term in pseudo_doc_terms:
368
- if term in tag_to_row_loaded:
369
- index = tag_to_row_loaded[term]
370
- pseudo_vector[index] = idf_loaded.get(term, 0)
371
-
372
- # Return the vector as a 2D array for compatibility with SVD transform
373
- return pseudo_vector.reshape(1, -1)
374
-
375
-
376
- def get_top_indices(reduced_pseudo_vector, reduced_matrix):
377
- # Compute cosine similarities
378
- similarities = cosine_similarity(reduced_pseudo_vector, reduced_matrix).flatten()
379
-
380
- # Get sorted tag indices based on similarities, in descending order
381
- sorted_indices = np.argsort(-similarities)
382
-
383
- # Return the top N indices
384
- return sorted_indices
385
-
386
-
387
- def get_tfidf_reduced_similar_tags(pseudo_doc_terms, allow_nsfw_tags):
388
- idf = tf_idf_components['idf']
389
- term_to_column_index = tf_idf_components['tag_to_column_index']
390
- row_to_tag = tf_idf_components['row_to_tag']
391
- reduced_matrix = tf_idf_components['reduced_matrix']
392
- svd = tf_idf_components['svd_model']
393
-
394
- # Construct the TF-IDF vector
395
- pseudo_tfidf_vector = construct_pseudo_vector(pseudo_doc_terms, idf, term_to_column_index)
396
-
397
- # Reduce the dimensionality of the pseudo-document vector for the reduced matrix
398
- reduced_pseudo_vector = svd.transform(pseudo_tfidf_vector)
399
-
400
- # Compute cosine similarities in the reduced space
401
- cosine_similarities_reduced = cosine_similarity(reduced_pseudo_vector, reduced_matrix).flatten()
402
-
403
- # Sort the indices by descending cosine similarity
404
- top_indices_reduced = np.argsort(cosine_similarities_reduced)
405
-
406
- # Map indices to tags with their similarities
407
- tag_similarity_dict = {row_to_tag[i]: cosine_similarities_reduced[i] for i in top_indices_reduced if i in row_to_tag}
408
-
409
- if not allow_nsfw_tags:
410
- tag_similarity_dict = {tag: sim for tag, sim in tag_similarity_dict.items() if tag not in nsfw_tags}
411
-
412
- tag_similarity_dict = {"by " + tag if is_artist(tag) else tag: sim for tag, sim in tag_similarity_dict.items()}
413
-
414
- # Sort and transform tag names
415
- sorted_tag_similarity_dict = OrderedDict(sorted(tag_similarity_dict.items(), key=lambda x: x[1], reverse=True))
416
- transformed_sorted_tag_similarity_dict = OrderedDict(
417
- (key.replace('_', ' ').replace('(', '\\(').replace(')', '\\)'), value)
418
- for key, value in sorted_tag_similarity_dict.items()
419
- )
420
-
421
- return transformed_sorted_tag_similarity_dict
422
-
423
-
424
- def create_html_placeholder(title="", content="", placeholder_height=400, placeholder_width="100%"):
425
- # Include a title in the same style as the top artists table heading
426
- html_placeholder = f"<div class=\"scrollable-content\" style='text-align: center;'><h1>{title}</h1></div>"
427
- # Conditionally add content if present
428
- if content:
429
- html_placeholder += f"<div style='text-align: center; margin-bottom: 20px;'><p>{content}</p></div>"
430
- # Add the placeholder div with specified height and width
431
- html_placeholder += f"<div style='height: {placeholder_height}px; width: {placeholder_width}; margin: 20px auto; background: transparent;'></div>"
432
- return html_placeholder
433
-
434
-
435
- def find_similar_tags(test_tags, tag_to_context_similarity, context_similarity_weight, allow_nsfw_tags):
436
- #Initialize stuff
437
- if not hasattr(find_similar_tags, "fasttext_small_model"):
438
- find_similar_tags.fasttext_small_model = compress_fasttext.models.CompressedFastTextKeyedVectors.load('e621FastTextModel010Replacement_small.bin')
439
- tag_aliases_file = 'fluffyrock_3m.csv'
440
- if not hasattr(find_similar_tags, "tag2aliases"):
441
- find_similar_tags.tag2aliases = build_aliases_dict(tag_aliases_file)
442
- if not hasattr(find_similar_tags, "alias2tags"):
443
- find_similar_tags.alias2tags = build_aliases_dict(tag_aliases_file, reverse=True)
444
- if not hasattr(find_similar_tags, "tag2count"):
445
- find_similar_tags.tag2count = build_tag_count_dict(tag_aliases_file)
446
- if not hasattr(find_similar_tags, "tag2idwiki"):
447
- find_similar_tags.tag2idwiki = build_tag_id_wiki_dict()
448
-
449
- modified_tags = [tag_info['modified_tag'] for tag_info in test_tags]
450
- transformed_tags = [tag.replace(' ', '_') for tag in modified_tags]
451
-
452
- # Find similar tags and prepare data for tables
453
- html_content = "<div class=\"scrollable-content\" style='display: inline-block; margin: 20px; text-align: center;'>"
454
- html_content += "<h1>Unknown Tags</h1>" # Heading for the table
455
- tags_added = False
456
- bad_entities = []
457
- known_entities_in_prompt = []
458
- encountered_modified_tags = set()
459
- for tag_info in test_tags:
460
- original_tag = tag_info['original_tag']
461
- modified_tag = tag_info['modified_tag']
462
- start_pos = tag_info['start_pos']
463
- end_pos = tag_info['end_pos']
464
- node_type = tag_info['node_type']
465
-
466
- if modified_tag in special_tags:
467
- bad_entities.append({"entity":"Special", "start":start_pos, "end":end_pos})
468
- continue
469
-
470
- if modified_tag in encountered_modified_tags:
471
- bad_entities.append({"entity":"Duplicate", "start":start_pos, "end":end_pos})
472
- continue
473
- encountered_modified_tags.add(modified_tag)
474
-
475
- modified_tag_for_search = modified_tag.replace(' ','_')
476
- similar_words = find_similar_tags.fasttext_small_model.most_similar(modified_tag_for_search, topn = 100)
477
- result, seen = [], set(transformed_tags)
478
-
479
- if modified_tag_for_search in find_similar_tags.tag2aliases:
480
- if modified_tag in find_similar_tags.tag2aliases and "_" in modified_tag: #Implicitly tell the user that they should get rid of the underscore
481
- result.append(modified_tag_for_search.replace('_',' '), 1)
482
- seen.add(modified_tag)
483
- else: #The user correctly did not put underscores in their tag
484
- count = find_similar_tags.tag2count.get(modified_tag_for_search, 0) # Get the count if available, otherwise default to 0
485
- tag_id, wiki_entry = find_similar_tags.tag2idwiki.get(modified_tag_for_search, (None, ''))
486
- # Check if tag_id and wiki_entry are valid
487
- wiki_url = ""
488
- if tag_id is not None and wiki_entry:
489
- # Construct the URL for the tag's wiki page
490
- wiki_url = f"https://e621.net/wiki_pages/{tag_id}"
491
- known_entities_in_prompt.append({"entity":"Known Tag", "start":start_pos, "end":end_pos, "count":count, "wiki_url":wiki_url, "wiki_entry":wiki_entry})
492
- continue
493
- else:
494
- for item in similar_words:
495
- similar_word, similarity = item
496
- if similar_word not in seen:
497
- if similar_word in find_similar_tags.tag2aliases:
498
- result.append((similar_word.replace('_', ' '), round(similarity, 3)))
499
- seen.add(similar_word)
500
- else:
501
- for similar_tag in find_similar_tags.alias2tags.get(similar_word, []):
502
- if similar_tag not in seen:
503
- result.append((similar_tag.replace('_', ' '), round(similarity, 3)))
504
- seen.add(similar_tag)
505
-
506
- #Remove NSFW tags if appropriate.
507
- if not allow_nsfw_tags:
508
- result = [(word, score) for word, score in result if word.replace(' ','_') not in nsfw_tags]
509
-
510
- #Adjust score based on context
511
- for i in range(len(result)):
512
- word, score = result[i] # Unpack the tuple
513
- context_score = tag_to_context_similarity.get(word,0)
514
- result[i] = (word, .5 * ((context_similarity_weight * context_score) + ((1 - context_similarity_weight) * score)))
515
-
516
- result = sorted(result, key=lambda x: x[1], reverse=True)[:10]
517
- html_content += create_html_tables_for_tags(modified_tag, "Corrected Tag", result, find_similar_tags.tag2count, find_similar_tags.tag2idwiki)
518
-
519
- bad_entities.append({"entity":"Unknown Tag", "start":start_pos, "end":end_pos})
520
-
521
- tags_added=True
522
- # If no tags were processed, add a message
523
- if not tags_added:
524
- html_content = create_html_placeholder(title="Unknown Tags", content="No Unknown Tags Found")
525
-
526
- return html_content, bad_entities, known_entities_in_prompt # Return list of lists for Dataframe
527
-
528
-
529
- def build_tag_offsets_dicts(new_image_tags_with_positions):
530
- # Structure the data for HighlightedText
531
- tag_data = []
532
- for tag_text, start_pos, nodetype in new_image_tags_with_positions:
533
- # Modify the tag
534
- modified_tag = tag_text.replace('_', ' ').replace('\\(', '(').replace('\\)', ')').strip()
535
- artist_matrix_tag = tag_text.replace('_', ' ').replace('\\(', '\(').replace('\\)', '\)').strip()
536
- tf_idf_matrix_tag = re.sub(r'\\([()])', r'\1', re.sub(r' ', '_', tag_text.strip().removeprefix('by ').removeprefix('by_')))
537
- # Calculate the end position based on the original tag length
538
- end_pos = start_pos + len(tag_text)
539
- # Append the structured data for each tag
540
- tag_data.append({
541
- "original_tag": tag_text,
542
- "start_pos": start_pos,
543
- "end_pos": end_pos,
544
- "modified_tag": modified_tag,
545
- "artist_matrix_tag": artist_matrix_tag,
546
- "tf_idf_matrix_tag": tf_idf_matrix_tag,
547
- "node_type": nodetype
548
- })
549
- return tag_data
550
-
551
-
552
- def augment_bad_entities_with_regex(text):
553
- bad_entities = []
554
-
555
- #comma at end
556
- match = re.search(r',(?=\s*$)', text)
557
- if match:
558
- index = match.start()
559
- bad_entities.append({"entity":"Remove Final Comma", "start":index, "end":index+1})
560
- match = re.search(r'\([^()]*(,)\s*\)\s*$', text)
561
- if match:
562
- index = match.start(1)
563
- bad_entities.append({"entity":"Remove Final Comma", "start":index, "end":index+1})
564
- match = re.search(r'\([^()]*(,)\s*:\s*\d+(\.\d+)?\s*\)\s*$', text)
565
- if match:
566
- index = match.start(1)
567
- bad_entities.append({"entity":"Remove Final Comma", "start":index, "end":index+1})
568
-
569
- # Comma after parentheses, multiple occurrences
570
- for match in re.finditer(r'(?<!\\)\)\s*(,)\s*[^\s]', text):
571
- index = match.start(1)
572
- bad_entities.append({"entity": "Move Comma Inside Parentheses", "start": index, "end": index + 1})
573
-
574
- # Double Comma detection
575
- for match in re.finditer(r',\s*,', text):
576
- index = match.start()
577
- bad_entities.append({"entity": "Double Comma", "start": index, "end": index + match.end() - match.start()})
578
-
579
- return bad_entities
580
-
581
- def escape_html(text):
582
- return text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;").replace('"', "&quot;").replace("'", "&#039;")
583
-
584
- def format_annotated_html(bad_entities, known_entities, text):
585
- tooltip_map = {
586
- "Unknown Tag": "This may not be a valid e621 tag. Consider removing or replacing it with tag(s) from the \"Unknown Tags\" section.",
587
- "Duplicate": "This tag has appeared multiple times in your prompt. Consider removing the copies.",
588
- "Remove Final Comma": "There should be no comma at the end of your prompt. Consider removing it.",
589
- "Move Comma Inside Parentheses": "In most e621-based models, the comma following a tag functions as an &quot;attention anchor&quot;, carrying most of the tag&apos;s information. It should therefore be assigned the same weight as the rest of the tag. So instead of &quot;(lineless:1.1),&quot;, consider &quot;(lineless,:1.1)&quot; or &quot;(lineless,)&quot;",
590
- "Double Comma": "One comma between tags is considered ample."
591
- }
592
- color_map = {
593
- "Unknown Tag": ("white", "red"), # White text on red background
594
- "Duplicate": ("black", "yellow"), # Black text on yellow background
595
- "Remove Final Comma": ("white", "blue"), # White text on blue background
596
- "Move Comma Inside Parentheses": ("white", "green"), # White text on green background
597
- "Double Comma": ("white","orange")
598
- }
599
-
600
- # Combine and sort entities
601
- combined_entities = bad_entities + known_entities
602
- combined_entities = sorted(combined_entities, key=lambda x: x['start'],reverse=True)
603
-
604
- # Generate HTML for the main text
605
- html_text = text
606
- for entity in combined_entities:
607
- start = entity['start']
608
- end = entity['end']
609
- label = entity['entity']
610
- if label == "Known Tag":
611
- wiki_url = entity.get('wiki_url', '')
612
- count = entity['count']
613
- wiki_entry = entity.get('wiki_entry', '')
614
- sanitized_wiki_entry = escape_html(wiki_entry) if wiki_entry else 'Unavailable'
615
- if wiki_url: # Check if wiki_url is not empty
616
- html_part = f'<a href="{wiki_url}" target="_blank" title="Count: {count}\tWiki: {sanitized_wiki_entry}" style="text-decoration: none; cursor: pointer; font-style: italic;">{text[start:end]}</a>'
617
- else:
618
- html_part = f'<span title="Count: {count}\tWiki: {sanitized_wiki_entry}" style="text-decoration: none; cursor: help; font-style: italic;">{text[start:end]}</span>'
619
- else:
620
- color = color_map.get(label, ("black", "white"))
621
- html_part = f'<span style="background-color: {color[1]}; color: {color[0]};">{text[start:end]}</span>'
622
- html_text = html_text[:start] + html_part + html_text[end:]
623
-
624
- # Generate HTML for the color key
625
- color_key_html = "<div style='text-align: right; margin-top: 20px;'>Key:"
626
- used_labels = set(entity['entity'] for entity in bad_entities)
627
- for label, colors in color_map.items():
628
- if label in used_labels:
629
- tooltip = tooltip_map.get(label, "")
630
- # Adding margin-right for spacing between items
631
- color_key_html += f" <span style='background-color: {colors[1]}; color: {colors[0]}; margin-right: 10px;' title='{tooltip}'>{label}</span>"
632
- color_key_html += "</div>"
633
-
634
- return f'<div style="padding: 10px; font-size: 16px;">{html_text}</div>{color_key_html}'
635
-
636
-
637
- def find_similar_artists(original_tags_string, top_n, context_similarity_weight, allow_nsfw_tags):
638
- try:
639
- new_tags_string = original_tags_string.lower()
640
- new_tags_string, removed_tags = remove_special_tags(new_tags_string)
641
-
642
- # Parse the prompt
643
- parsed = parser.parse(new_tags_string)
644
- # Extract tags from the parsed tree
645
- new_image_tags = extract_tags(parsed)
646
- tag_data = build_tag_offsets_dicts(new_image_tags)
647
-
648
- #Suggested tags stuff
649
- suggested_tags_html_content = "<div class=\"scrollable-content\" style='display: inline-block; margin: 20px; text-align: center;'>"
650
- suggested_tags_html_content += "<h1>Suggested Tags</h1>" # Heading for the table
651
- suggested_tags = get_tfidf_reduced_similar_tags([item["tf_idf_matrix_tag"] for item in tag_data] + removed_tags, allow_nsfw_tags)
652
-
653
- unseen_tags_data, bad_entities, known_entities = find_similar_tags(tag_data, suggested_tags, context_similarity_weight, allow_nsfw_tags)
654
-
655
- #Bad tags stuff
656
- bad_entities.extend(augment_bad_entities_with_regex(new_tags_string))
657
- bad_entities.sort(key=lambda x: x['start'])
658
- #bad_tags_illustrated_string = {"text":new_tags_string, "entities":bad_entities}
659
- bad_tags_illustrated_html = format_annotated_html(bad_entities, known_entities, new_tags_string)
660
-
661
- # Create a set of tags that should be filtered out
662
- filter_tags = {entry["original_tag"].strip() for entry in tag_data}
663
- # Use this set to filter suggested_tags
664
- suggested_tags_filtered = OrderedDict((k, v) for k, v in suggested_tags.items() if k not in filter_tags)
665
-
666
- # Splitting the dictionary into two based on the condition
667
- suggested_artist_tags_filtered = OrderedDict((k, v) for k, v in suggested_tags_filtered.items() if k.startswith("by "))
668
- suggested_non_artist_tags_filtered = OrderedDict((k, v) for k, v in suggested_tags_filtered.items() if not k.startswith("by ") and k not in special_tags)
669
-
670
- topnsuggestions = list(islice(suggested_non_artist_tags_filtered.items(), 100))
671
- suggested_tags_html_content += create_html_tables_for_tags("-", "Suggested Tag", topnsuggestions, find_similar_tags.tag2count, find_similar_tags.tag2idwiki)
672
-
673
- #Artist stuff
674
- excluded_artists = ["by conditional dnp", "by unknown artist"]
675
- top_artists = [(key, value) for key, value in suggested_artist_tags_filtered.items() if key.lower() not in excluded_artists][:top_n]
676
- top_artists_str = create_top_artists_table(top_artists)
677
- dynamic_prompts_formatted_artists = "{" + "|".join([artist for artist, _ in top_artists]) + "}"
678
-
679
- image_galleries = []
680
- for root, dirs, files in os.walk(sample_images_directory_path):
681
- for name in dirs:
682
- baseline, artists = generate_artist_image_tuples([name[3:] for name, _ in top_artists], os.path.join(root, name))
683
- image_galleries.append(baseline) # Add baseline as its own gallery item
684
- image_galleries.append(artists) # Extend the list with artist tuples
685
-
686
- return (unseen_tags_data, bad_tags_illustrated_html, suggested_tags_html_content, top_artists_str, dynamic_prompts_formatted_artists, *image_galleries)
687
- except ParseError as e:
688
- return [], "Parse Error: Check for mismatched parentheses or something", "", "", None, None
689
-
690
-
691
- with gr.Blocks(css=css) as app:
692
- with gr.Group():
693
- with gr.Row():
694
- with gr.Column(scale=3):
695
- image_tags = gr.Textbox(label="Enter Prompt", placeholder="e.g. fox, outside, detailed background, ...")
696
- #bad_tags_illustrated_string = gr.HighlightedText(show_legend=True, color_map={"Unknown Tag":"red","Duplicate":"yellow","Remove Final Comma":"purple","Move Comma Inside Parentheses":"green"}, label="Annotated Prompt")
697
- bad_tags_illustrated_string = gr.HTML()
698
- with gr.Column(scale=1):
699
- #image_path = os.path.join("https://huggingface.co/spaces/FoodDesert/Prompt_Squirrel/resolve/main", "transparentsquirrel.png")
700
- #gr.Image(label=" ", value=image_path, height=155, width=140)
701
- gr.HTML('<div style="text-align: center;"><img src="https://huggingface.co/spaces/FoodDesert/Prompt_Squirrel/resolve/main/mascotimages/transparentsquirrel.png" alt="Cute Mascot" style="height: 220px; width: auto; background: transparent;"></div><br>')
702
- #gr.HTML("<br>" * 2) # Adjust the number of line breaks ("<br>") as needed to push the button down
703
- #image_path = os.path.join('mascotimages', "transparentsquirrel.png")
704
- #random_image_path = os.path.join('mascotimages', random.choice([f for f in os.listdir('mascotimages') if os.path.isfile(os.path.join('mascotimages', f))]))
705
- #with Image.open(random_image_path) as img:
706
- # gr.Image(value=img,show_label=False, show_download_button=False, show_share_button=False, height=200)
707
- #gr.Image(value="https://huggingface.co/spaces/FoodDesert/Prompt_Squirrel/resolve/main/mascotimages/transparentsquirrel.png",show_label=False, show_download_button=False, show_share_button=False, height=200)
708
- #I posted the image to discord, and that's where this link came from. This is a very ugly way to do this, but I could not, no matter what I tried, get it to display an image from within the space itself. The galleries work fine for some reason, but not this.
709
- #gr.Image(value="https://res.cloudinary.com/dnse84ol6/image/upload/v1713538125/transparentsquirrel_zhou7f.png",show_label=False, show_download_button=False, show_share_button=False, height=200)
710
- submit_button = gr.Button(variant="primary")
711
- with gr.Row():
712
- with gr.Column(scale=3):
713
- with gr.Group():
714
- with gr.Row():
715
- context_similarity_weight = gr.Slider(minimum=0, maximum=1, value=0.5, step=0.1, label="Context Similarity Weight")
716
- allow_nsfw = gr.Checkbox(label="Allow NSFW Tags", value=False)
717
- with gr.Row():
718
- with gr.Column(scale=2):
719
- unseen_tags = gr.HTML(label="Unknown Tags", value=create_html_placeholder(title="Unknown Tags"))
720
- with gr.Column(scale=1):
721
- suggested_tags = gr.HTML(label="Suggested Tags", value=create_html_placeholder(title="Suggested Tags"))
722
- with gr.Column(scale=1):
723
- with gr.Group():
724
- num_artists = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Number of artists")
725
- top_artists = gr.HTML(label="Top Artists", value=create_html_placeholder(title="Top Artists"))
726
- dynamic_prompts = gr.Textbox(label="Dynamic Prompts Format", info="For if you're using the Automatic1111 webui (https://github.com/AUTOMATIC1111/stable-diffusion-webui) with the Dynamic Prompts extension activated (https://github.com/adieyal/sd-dynamic-prompts) and want to try them all individually.")
727
- galleries = []
728
- for root, dirs, files in os.walk(sample_images_directory_path):
729
- for name in dirs:
730
- with gr.Row():
731
- baseline = gr.Gallery(allow_preview=False, rows=1, columns=1, height=420, scale=3)
732
- styles = gr.Gallery(preview=False, rows=2, columns=5, height=420, scale=8)
733
- galleries.extend([baseline, styles])
734
-
735
- submit_button.click(
736
- find_similar_artists,
737
- inputs=[image_tags, num_artists, context_similarity_weight, allow_nsfw],
738
- outputs=[unseen_tags, bad_tags_illustrated_string, suggested_tags, top_artists, dynamic_prompts] + galleries
739
- )
740
-
741
- gr.Markdown(faq_content)
742
-
743
-
744
- app.launch()
745
-
 
1
+ import gradio as gr
2
+ from sklearn.metrics.pairwise import cosine_similarity
3
+ from scipy.sparse import csr_matrix
4
+ import numpy as np
5
+ import joblib
6
+ from joblib import load
7
+ import h5py
8
+ from io import BytesIO
9
+ import csv
10
+ import re
11
+ import random
12
+ import compress_fasttext
13
+ from collections import OrderedDict
14
+ from lark import Lark, Tree, Token
15
+ from lark.exceptions import ParseError
16
+ import json
17
+ import zipfile
18
+ from PIL import Image
19
+ import io
20
+ import os
21
+ import glob
22
+ import itertools
23
+ from itertools import islice
24
+ from pathlib import Path
25
+ import logging
26
+
27
+ # Set up logging
28
+ logging.basicConfig(filename='error.log', level=logging.DEBUG, format='%(asctime)s %(levelname)s:%(message)s')
29
+
30
+
31
+ faq_content="""
32
+ # Questions:
33
+
34
+ ## What is the purpose of this tool?
35
+
36
+ Since Stable Diffusion's initial release in 2022, users have developed a myriad of fine-tuned text to image models, each with unique "linguistic" preferences depending on the data from which it was fine-tuned.
37
+ Some models react best when prompted with verbose scene descriptions akin to DALL-E, while others fine-tuned on images scraped from popular image boards understand those boards' tag sets.
38
+ This tool serves as a linguistic bridge to the e621 image board tag lexicon, on which many popular models such as Fluffyrock, Fluffusion, and Pony Diffusion v6 were trained.
39
+
40
+ When you enter a txt2img prompt and press the "submit" button, Prompt Squirrel parses your prompt and checks that all your tags are valid e621 tags.
41
+ If it finds any that are not, it recommends some valid e621 tags you can use to replace them in the "Unknown Tags" section.
42
+ Additionally, in the "Top Artists" text box, it lists the artists who would most likely draw an image having the set of tags you provided.
43
+ This is useful to align your prompt with the expected input to an e621-trained model.
44
+
45
+ ## Does input order matter?
46
+
47
+ No
48
+
49
+ ## Should I use underscores or spaces in the input tags?
50
+
51
+ As a rule, e621-trained models replace underscores in tags with spaces, so spaces are preferred.
52
+
53
+ ## Can I use parentheses or weights as in the Stable Diffusion Automatic1111 WebUI?
54
+
55
+ Yes, but only '(' and ')' and numerical weights, and all of these things are ignored in all calculations. The main benefit of this is that you can copy/paste prompts from one program to another with minimal editing.
56
+ An example that illustrates acceptable parentheses and weight formatting is:
57
+ ((sunset over the mountains)), (clear sky:1.5), ((eagle flying high:2.0)), river, (fish swimming in the river:1.2), (campfire, (marshmallows:2.1):1.3), stars in the sky, ((full moon:1.8)), (wolf howling:1.7)
58
+
59
+ ## Why are some valid tags marked as "unknown", and why don't some artists ever get returned?
60
+
61
+ Some data is excluded from consideration if it did not occur frequently enough in the sample from which the application makes its calculations.
62
+ If an artist or tag is too infrequent, we might not think we have enough data to make predictions about it. Additionally, Prompt Squirrel gathers information from several sources, and the sources are not always consistent about things like exact tag names or counts, which vary over time.
63
+
64
+ ## Why do some suggested tags not have summaries or wiki links, and of those that do, why do some look truncated?
65
+
66
+ Both of these features are extracted from the tag wiki pages, but some valid e621 tags do not have wiki pages. Additionally, the summaries are heuristically extracted from the beginning of the wiki pages, and this extraction process is prone to some amount of error.
67
+
68
+ ## Are there any special tags?
69
+
70
+ Yes. We normalized the favorite counts of each image to a range of 0-9, with 0 being the lowest favcount, and 9 being the highest.
71
+ You can include any of these special tags: "score:0", "score:1", "score:2", "score:3", "score:4", "score:5", "score:6", "score:7", "score:8", "score:9"
72
+ in your list to bias the output toward artists with higher or lower scoring images.
73
+
74
+ ## Are there any other special tricks?
75
+
76
+ Yes. If you want to more strongly bias the artist output toward a specific tag, you can just list it multiple times.
77
+ So for example, the query "red fox, red fox, red fox, score:7" will yield a list of artists who are more strongly associated with the tag "red fox"
78
+ than the query "red fox, score:7".
79
+
80
+ ## Why is this space tagged "not-for-all-audience"
81
+ The "not-for-all-audience" tag informs users that this tool's text output is derived from e621.net data for tag prediction and completion.
82
+ The app will try not to display nsfw tags unless the "Allow NSFW Tags" is checked, but the filter is not perfect.
83
+
84
+ ## How is the artist list calculated?
85
+
86
+ Each artist is represented by a "pseudo-document" composed of all the tags from their uploaded images, treating these tags similarly to words in a text document.
87
+ Similarly, when you input a set of tags, the system creates a pseudo-document for your query out of all the tags.
88
+ It then compares your tags against each artist's collection, essentially finding which artist's tags are most "similar" to yours.
89
+ This method helps identify artists whose work is closely aligned with the themes or elements you're interested in.
90
+ For those curious about the underlying mechanics of comparing text-like data, we employ the TF-IDF (Term Frequency-Inverse Document Frequency) method, a standard approach in information retrieval, and reduce the TF-IDF matrix to a reasonable size using Singular Value Decomposition.
91
+ You can read more about TF-IDF on its [Wikipedia page](https://en.wikipedia.org/wiki/Tf%E2%80%93idf) and Singular Value Decomposition on its [Wikipedia page](https://en.wikipedia.org/wiki/Singular_value_decomposition).
92
+
93
+ ## How does the tag corrector work?
94
+
95
+ We collect the tag sets from over 4 million e621 posts, treating the tag set from each image as an individual document.
96
+ We then randomly replace about 10% of the tags in each document with a randomly selected alias from e621's list of aliases for the tag
97
+ (e.g. "canine" gets replaced with one of {k9,canines,mongrel,cannine,cnaine,feral_canine,anthro_canine}).
98
+ We then train a FastText (https://fasttext.cc/) model on the documents. The result of this training is a function that maps arbitrary words to vectors such that
99
+ the vector for a tag and the vectors for its aliases are all close together (because the model has seen them in similar contexts).
100
+ Since the lists of aliases contain misspellings and rephrasings of tags, the model should be robust to these kinds of problems as long as they are not too dissimilar from the alias lists.
101
+
102
+ To enhance the tag corrector further, we employ the same TF-IDF method we used for artist tags to calculate a separate, context-sensitive similarity score for each of the top 100 tags selected via the FastText method.
103
+ By considering the context in which tags are used, we can now not only correct misspellings and rephrasings but also make more contextually relevant suggestions.
104
+ The "similarity weight" slider controls how much weight these TF-IDF scores are given vs how much weight the FastText similarity model is given when suggesting replacements for invalid tags.
105
+ A similarity weight slider value of 0 means that only the FastText model's predictions will be used to calculate similarity scores, and a value of 1 means only the TF-IDF scores are used (although the FastText model is still used to trim the list of candidates).
106
+
107
+
108
+ ## How do the sample images work?
109
+
110
+ In the first row of galleries, for each artist in the dataset, we generated a sample image with the model Fluffyrock Unleashed using the prompt "by artist, soyjak, anthro, male, bust portrait, meme, grin" where "artist" is the name of an artist.
111
+ The simplicity of the prompt, the the simplicty of the default style, and the recognizability of the character make it easier to understand how artist names affect generated image styles.
112
+ The image on the left captioned "No Artist" was generated with the same prompt, but with no artist name.
113
+ You should compare all the images to the first to see how the artist names affect the output.
114
+ Each subsequent row of images was generated using the same process, but with a different prompt.
115
+ See SamplePrompts.csv for the list of prompts used and their descriptions.
116
+ """
117
+
118
+
119
+ nsfw_threshold = 0.95 # Assuming the threshold value is defined here
120
+
121
+ css = """
122
+ .scrollable-content {
123
+ max-height: 500px;
124
+ overflow-y: auto;
125
+ }
126
+ """
127
+
128
+ grammar=r"""
129
+ !start: (prompt | /[][():]/+)*
130
+ prompt: (emphasized | plain | comma | WHITESPACE)*
131
+ !emphasized: "(" prompt ")"
132
+ | "(" prompt ":" [WHITESPACE] NUMBER [WHITESPACE] ")"
133
+ comma: ","
134
+ WHITESPACE: /\s+/
135
+ plain: /([^,\\\[\]():|]|\\.)+/
136
+ %import common.SIGNED_NUMBER -> NUMBER
137
+ """
138
+
139
+ # Initialize the parser
140
+ parser = Lark(grammar, start='start')
141
+
142
+ # Function to extract tags
143
+ def extract_tags(tree):
144
+ tags_with_positions = []
145
+ def _traverse(node):
146
+ if isinstance(node, Token) and node.type == '__ANON_1':
147
+ tag_position = node.start_pos
148
+ tag_text = node.value
149
+ tags_with_positions.append((tag_text, tag_position, "tag"))
150
+ elif not isinstance(node, Token):
151
+ for child in node.children:
152
+ _traverse(child)
153
+ _traverse(tree)
154
+ return tags_with_positions
155
+
156
+
157
+ special_tags = ["score:0", "score:1", "score:2", "score:3", "score:4", "score:5", "score:6", "score:7", "score:8", "score:9", "rating:s", "rating:q", "rating:e"]
158
+ def remove_special_tags(original_string):
159
+ tags = [tag.strip() for tag in original_string.split(",")]
160
+ remaining_tags = [tag for tag in tags if tag not in special_tags]
161
+ removed_tags = [tag for tag in tags if tag in special_tags]
162
+ return ", ".join(remaining_tags), removed_tags
163
+
164
+
165
+ # Define a function to load all necessary components
166
+ def load_model_components(file_path):
167
+ # Ensure the file path is a Path object for robust path handling
168
+ file_path = Path(file_path)
169
+
170
+ # Check if the file exists
171
+ if not file_path.is_file():
172
+ raise FileNotFoundError(f"The specified joblib file was not found: {file_path}")
173
+
174
+ # Load all the model components from the joblib file
175
+ model_components = joblib.load(file_path)
176
+
177
+ # Create a reverse mapping from row index to tag
178
+ if 'tag_to_row_index' in model_components:
179
+ model_components['row_to_tag'] = {idx: tag for tag, idx in model_components['tag_to_row_index'].items()}
180
+
181
+ return model_components
182
+
183
+ # Load all components at the start
184
+ tf_idf_components = load_model_components('tf_idf_files_420.joblib')
185
+
186
+
187
+ nsfw_tags = set() # Initialize an empty set to store words meeting the threshold
188
+ # Open and read the CSV file
189
+ with open("word_rating_probabilities.csv", 'r', newline='', encoding='utf-8') as csvfile:
190
+ reader = csv.reader(csvfile)
191
+ next(reader, None) # Skip the header row
192
+ for row in reader:
193
+ word = row[0] # The word is in the first column
194
+ probability_sum = float(row[1]) # The sum of probabilities is in the second column, convert to float for comparison
195
+ # Check if the probability sum meets the threshold and add the word to the set if it does
196
+ if probability_sum >= nsfw_threshold:
197
+ nsfw_tags.add(word)
198
+
199
+
200
+ # Read the set of valid artists into memory.
201
+ artist_set = set()
202
+ with open("fluffyrock_3m.csv", 'r', newline='', encoding='utf-8') as csvfile:
203
+ """
204
+ Load artist names from a CSV file and store them in the global set.
205
+ Artist tags start with 'by_' and the prefix will be removed.
206
+ """
207
+ reader = csv.reader(csvfile)
208
+ for row in reader:
209
+ tag_name = row[0] # Assuming the first column contains the tag names
210
+ if tag_name.startswith('by_'):
211
+ # Strip 'by_' from the start of the tag name and add to the set
212
+ artist_name = tag_name[3:] # Remove the first three characters 'by_'
213
+ artist_set.add(artist_name)
214
+ def is_artist(name):
215
+ return name in artist_set
216
+
217
+
218
+ sample_images_directory_path = 'sampleimages'
219
+ def generate_artist_image_tuples(top_artists, image_directory):
220
+ json_files = glob.glob(f'{image_directory}/*.json')
221
+ json_file_path = json_files[0] if json_files else None
222
+ with open(json_file_path, 'r') as json_file:
223
+ artist_to_file_map = json.load(json_file)
224
+
225
+ filename = artist_to_file_map.get("")
226
+ image_path = os.path.join(image_directory, filename)
227
+ if os.path.exists(image_path):
228
+ baseline_tuple = [(image_path, "No Artist")]
229
+
230
+ artist_image_tuples = []
231
+ for artist in top_artists:
232
+ filename = artist_to_file_map.get(artist)
233
+ if filename:
234
+ image_path = os.path.join(image_directory, filename)
235
+ if os.path.exists(image_path):
236
+ artist_image_tuples.append((image_path, artist if artist else "No Artist"))
237
+
238
+ return baseline_tuple, artist_image_tuples
239
+
240
+
241
+ def clean_tag(tag):
242
+ return ''.join(char for char in tag if ord(char) < 128)
243
+
244
+
245
+ #Normally returns tag to aliases, but when reverse=True, returns alias to tags
246
+ def build_aliases_dict(filename, reverse=False):
247
+ aliases_dict = {}
248
+ with open(filename, 'r', newline='', encoding='utf-8') as csvfile:
249
+ reader = csv.reader(csvfile)
250
+ for row in reader:
251
+ tag = clean_tag(row[0])
252
+ alias_list = [] if row[3] == "null" else [clean_tag(alias) for alias in row[3].split(',')]
253
+ if reverse:
254
+ for alias in alias_list:
255
+ aliases_dict.setdefault(alias, []).append(tag)
256
+ else:
257
+ aliases_dict[tag] = alias_list
258
+ return aliases_dict
259
+
260
+
261
+ def build_tag_count_dict(filename):
262
+ with open(filename, 'r', newline='', encoding='utf-8') as csvfile:
263
+ reader = csv.reader(csvfile)
264
+ result_dict = {}
265
+ for row in reader:
266
+ key = row[0]
267
+ value = int(row[2]) if row[2].isdigit() else None
268
+ if value is not None:
269
+ result_dict[key] = value
270
+ return result_dict
271
+
272
+ import csv
273
+
274
+
275
+ def build_tag_id_wiki_dict(filename='wiki_pages-2023-08-08.csv'):
276
+ """
277
+ Reads a CSV file and returns a dictionary mapping tag names to tuples of
278
+ (number, most relevant line from the wiki entry). Rows with a non-integer in the first column are ignored.
279
+ The most relevant line is the first line that does not start with "thumb" and is not blank.
280
+
281
+ Parameters:
282
+ - filename: The path to the CSV file.
283
+
284
+ Returns:
285
+ - A dictionary where each key is a tag name and each value is a tuple (number, most relevant wiki entry line).
286
+ """
287
+ tag_data = {}
288
+ with open(filename, 'r', encoding='utf-8') as csvfile:
289
+ reader = csv.reader(csvfile)
290
+
291
+ # Skip the header row
292
+ next(reader)
293
+
294
+ for row in reader:
295
+ try:
296
+ # Attempt to convert the first column to an integer
297
+ number = int(row[0])
298
+ except ValueError:
299
+ # If conversion fails, skip this row
300
+ continue
301
+
302
+ tag = row[3]
303
+ wiki_entry_full = row[4]
304
+
305
+ # Process the wiki_entry to find the most relevant line
306
+ relevant_line = ''
307
+ for line in wiki_entry_full.split('\n'):
308
+ if line.strip() and not line.startswith("thumb"):
309
+ relevant_line = line
310
+ break
311
+
312
+ # Map the tag to a tuple of (number, relevant_line)
313
+ tag_data[tag] = (number, relevant_line)
314
+
315
+ return tag_data
316
+
317
+
318
+ def create_html_tables_for_tags(subtable_heading, item_heading, word_similarity_tuples, tag2count, tag2idwiki):
319
+ # Wrap the tag part in a <span> with styles for bold and larger font
320
+ html_str = f"<div style='display: inline-block; margin: 20px; vertical-align: top;'><table><thead><tr><th colspan='3' style='text-align: center; padding-bottom: 10px;'><span style='font-weight: bold; font-size: 20px;'>{subtable_heading}</span></th></tr></thead><tbody><tr style='border-bottom: 1px solid #000;'><th>{item_heading}</th><th>Similarity</th><th>Count</th></tr>"
321
+ # Loop through the results and add table rows for each
322
+ for word, sim in word_similarity_tuples:
323
+ word_with_underscores = word.replace(' ', '_')
324
+ word_with_escaped_parentheses = word.replace("\\(", "(").replace("\\)", ")").replace("(", "\\(").replace(")", "\\)")
325
+ count = tag2count.get(word_with_underscores.replace("\\(", "(").replace("\\)", ")"), 0) # Get the count if available, otherwise default to 0
326
+ tag_id, wiki_entry = tag2idwiki.get(word_with_underscores, (None, ''))
327
+ # Check if tag_id and wiki_entry are valid
328
+ if tag_id is not None and wiki_entry:
329
+ # Construct the URL for the tag's wiki page
330
+ wiki_url = f"https://e621.net/wiki_pages/{tag_id}"
331
+ # Make the tag a hyperlink with a tooltip
332
+ tag_element = f"<a href='{wiki_url}' target='_blank' title='{wiki_entry}'>{word_with_escaped_parentheses}</a>"
333
+ else:
334
+ # Display the word without any hyperlink or tooltip
335
+ tag_element = word_with_escaped_parentheses
336
+ # Include the tag element in the table row
337
+ html_str += f"<tr><td style='border: none; padding: 5px; height: 20px;'>{tag_element}</td><td style='border: none; padding: 5px; height: 20px;'>{round(sim, 3)}</td><td style='border: none; padding: 5px; height: 20px;'>{count}</td></tr>"
338
+
339
+ html_str += "</tbody></table></div>"
340
+ return html_str
341
+
342
+
343
+ def create_top_artists_table(top_artists):
344
+ # Add a heading above the table
345
+ html_str = "<div class=\"scrollable-content\" style='display: inline-block; margin: 20px; text-align: center;'>"
346
+ html_str += "<h1>Top Artists</h1>" # Heading for the table
347
+ # Start the table with increased font size and no borders between rows
348
+ html_str += "<table style='font-size: 20px; border-collapse: collapse;'>"
349
+ html_str += "<thead><tr><th>Artist</th><th>Similarity</th></tr></thead><tbody>"
350
+ # Loop through the top artists and add a row for each without the rank and without borders between rows
351
+ for artist, score in top_artists:
352
+ artist_name = artist[3:] if artist.startswith("by ") else artist # Remove "by " prefix
353
+ similarity_percentage = "{:.1f}%".format(score * 100) # Convert score to percentage string with one decimal
354
+ html_str += f"<td style='padding: 3px 20px; border: none;'>{artist_name}</td><td style='padding: 3px 20px; border: none;'>{similarity_percentage}</td></tr>"
355
+
356
+ # Close the table HTML
357
+ html_str += "</tbody></table></div>"
358
+
359
+ return html_str
360
+
361
+
362
+ def construct_pseudo_vector(pseudo_doc_terms, idf_loaded, tag_to_row_loaded):
363
+ # Initialize a vector of zeros with the length of the term_to_index mapping
364
+ pseudo_vector = np.zeros(len(tag_to_row_loaded))
365
+
366
+ # Fill in the vector for terms in the pseudo document
367
+ for term in pseudo_doc_terms:
368
+ if term in tag_to_row_loaded:
369
+ index = tag_to_row_loaded[term]
370
+ pseudo_vector[index] = idf_loaded.get(term, 0)
371
+
372
+ # Return the vector as a 2D array for compatibility with SVD transform
373
+ return pseudo_vector.reshape(1, -1)
374
+
375
+
376
+ def get_top_indices(reduced_pseudo_vector, reduced_matrix):
377
+ # Compute cosine similarities
378
+ similarities = cosine_similarity(reduced_pseudo_vector, reduced_matrix).flatten()
379
+
380
+ # Get sorted tag indices based on similarities, in descending order
381
+ sorted_indices = np.argsort(-similarities)
382
+
383
+ # Return the top N indices
384
+ return sorted_indices
385
+
386
+
387
+ def get_tfidf_reduced_similar_tags(pseudo_doc_terms, allow_nsfw_tags):
388
+ idf = tf_idf_components['idf']
389
+ term_to_column_index = tf_idf_components['tag_to_column_index']
390
+ row_to_tag = tf_idf_components['row_to_tag']
391
+ reduced_matrix = tf_idf_components['reduced_matrix']
392
+ svd = tf_idf_components['svd_model']
393
+
394
+ # Construct the TF-IDF vector
395
+ pseudo_tfidf_vector = construct_pseudo_vector(pseudo_doc_terms, idf, term_to_column_index)
396
+
397
+ # Reduce the dimensionality of the pseudo-document vector for the reduced matrix
398
+ reduced_pseudo_vector = svd.transform(pseudo_tfidf_vector)
399
+
400
+ # Compute cosine similarities in the reduced space
401
+ cosine_similarities_reduced = cosine_similarity(reduced_pseudo_vector, reduced_matrix).flatten()
402
+
403
+ # Sort the indices by descending cosine similarity
404
+ top_indices_reduced = np.argsort(cosine_similarities_reduced)
405
+
406
+ # Map indices to tags with their similarities
407
+ tag_similarity_dict = {row_to_tag[i]: cosine_similarities_reduced[i] for i in top_indices_reduced if i in row_to_tag}
408
+
409
+ if not allow_nsfw_tags:
410
+ tag_similarity_dict = {tag: sim for tag, sim in tag_similarity_dict.items() if tag not in nsfw_tags}
411
+
412
+ tag_similarity_dict = {"by " + tag if is_artist(tag) else tag: sim for tag, sim in tag_similarity_dict.items()}
413
+
414
+ # Sort and transform tag names
415
+ sorted_tag_similarity_dict = OrderedDict(sorted(tag_similarity_dict.items(), key=lambda x: x[1], reverse=True))
416
+ transformed_sorted_tag_similarity_dict = OrderedDict(
417
+ (key.replace('_', ' ').replace('(', '\\(').replace(')', '\\)'), value)
418
+ for key, value in sorted_tag_similarity_dict.items()
419
+ )
420
+
421
+ return transformed_sorted_tag_similarity_dict
422
+
423
+
424
+ def create_html_placeholder(title="", content="", placeholder_height=400, placeholder_width="100%"):
425
+ # Include a title in the same style as the top artists table heading
426
+ html_placeholder = f"<div class=\"scrollable-content\" style='text-align: center;'><h1>{title}</h1></div>"
427
+ # Conditionally add content if present
428
+ if content:
429
+ html_placeholder += f"<div style='text-align: center; margin-bottom: 20px;'><p>{content}</p></div>"
430
+ # Add the placeholder div with specified height and width
431
+ html_placeholder += f"<div style='height: {placeholder_height}px; width: {placeholder_width}; margin: 20px auto; background: transparent;'></div>"
432
+ return html_placeholder
433
+
434
+
435
+ def find_similar_tags(test_tags, tag_to_context_similarity, context_similarity_weight, allow_nsfw_tags):
436
+ #Initialize stuff
437
+ if not hasattr(find_similar_tags, "fasttext_small_model"):
438
+ find_similar_tags.fasttext_small_model = compress_fasttext.models.CompressedFastTextKeyedVectors.load('e621FastTextModel010Replacement_small.bin')
439
+ tag_aliases_file = 'fluffyrock_3m.csv'
440
+ if not hasattr(find_similar_tags, "tag2aliases"):
441
+ find_similar_tags.tag2aliases = build_aliases_dict(tag_aliases_file)
442
+ if not hasattr(find_similar_tags, "alias2tags"):
443
+ find_similar_tags.alias2tags = build_aliases_dict(tag_aliases_file, reverse=True)
444
+ if not hasattr(find_similar_tags, "tag2count"):
445
+ find_similar_tags.tag2count = build_tag_count_dict(tag_aliases_file)
446
+ if not hasattr(find_similar_tags, "tag2idwiki"):
447
+ find_similar_tags.tag2idwiki = build_tag_id_wiki_dict()
448
+
449
+ modified_tags = [tag_info['modified_tag'] for tag_info in test_tags]
450
+ transformed_tags = [tag.replace(' ', '_') for tag in modified_tags]
451
+
452
+ # Find similar tags and prepare data for tables
453
+ html_content = "<div class=\"scrollable-content\" style='display: inline-block; margin: 20px; text-align: center;'>"
454
+ html_content += "<h1>Unknown Tags</h1>" # Heading for the table
455
+ tags_added = False
456
+ bad_entities = []
457
+ known_entities_in_prompt = []
458
+ encountered_modified_tags = set()
459
+ for tag_info in test_tags:
460
+ original_tag = tag_info['original_tag']
461
+ modified_tag = tag_info['modified_tag']
462
+ start_pos = tag_info['start_pos']
463
+ end_pos = tag_info['end_pos']
464
+ node_type = tag_info['node_type']
465
+
466
+ if modified_tag in special_tags:
467
+ bad_entities.append({"entity":"Special", "start":start_pos, "end":end_pos})
468
+ continue
469
+
470
+ if modified_tag in encountered_modified_tags:
471
+ bad_entities.append({"entity":"Duplicate", "start":start_pos, "end":end_pos})
472
+ continue
473
+ encountered_modified_tags.add(modified_tag)
474
+
475
+ modified_tag_for_search = modified_tag.replace(' ','_')
476
+ similar_words = find_similar_tags.fasttext_small_model.most_similar(modified_tag_for_search, topn = 100)
477
+ result, seen = [], set(transformed_tags)
478
+
479
+ if modified_tag_for_search in find_similar_tags.tag2aliases:
480
+ if modified_tag in find_similar_tags.tag2aliases and "_" in modified_tag: #Implicitly tell the user that they should get rid of the underscore
481
+ result.append(modified_tag_for_search.replace('_',' '), 1)
482
+ seen.add(modified_tag)
483
+ else: #The user correctly did not put underscores in their tag
484
+ count = find_similar_tags.tag2count.get(modified_tag_for_search, 0) # Get the count if available, otherwise default to 0
485
+ tag_id, wiki_entry = find_similar_tags.tag2idwiki.get(modified_tag_for_search, (None, ''))
486
+ # Check if tag_id and wiki_entry are valid
487
+ wiki_url = ""
488
+ if tag_id is not None and wiki_entry:
489
+ # Construct the URL for the tag's wiki page
490
+ wiki_url = f"https://e621.net/wiki_pages/{tag_id}"
491
+ known_entities_in_prompt.append({"entity":"Known Tag", "start":start_pos, "end":end_pos, "count":count, "wiki_url":wiki_url, "wiki_entry":wiki_entry})
492
+ continue
493
+ else:
494
+ for item in similar_words:
495
+ similar_word, similarity = item
496
+ if similar_word not in seen:
497
+ if similar_word in find_similar_tags.tag2aliases:
498
+ result.append((similar_word.replace('_', ' '), round(similarity, 3)))
499
+ seen.add(similar_word)
500
+ else:
501
+ for similar_tag in find_similar_tags.alias2tags.get(similar_word, []):
502
+ if similar_tag not in seen:
503
+ result.append((similar_tag.replace('_', ' '), round(similarity, 3)))
504
+ seen.add(similar_tag)
505
+
506
+ #Remove NSFW tags if appropriate.
507
+ if not allow_nsfw_tags:
508
+ result = [(word, score) for word, score in result if word.replace(' ','_') not in nsfw_tags]
509
+
510
+ #Adjust score based on context
511
+ for i in range(len(result)):
512
+ word, score = result[i] # Unpack the tuple
513
+ context_score = tag_to_context_similarity.get(word,0)
514
+ result[i] = (word, .5 * ((context_similarity_weight * context_score) + ((1 - context_similarity_weight) * score)))
515
+
516
+ result = sorted(result, key=lambda x: x[1], reverse=True)[:10]
517
+ html_content += create_html_tables_for_tags(modified_tag, "Corrected Tag", result, find_similar_tags.tag2count, find_similar_tags.tag2idwiki)
518
+
519
+ bad_entities.append({"entity":"Unknown Tag", "start":start_pos, "end":end_pos})
520
+
521
+ tags_added=True
522
+ # If no tags were processed, add a message
523
+ if not tags_added:
524
+ html_content = create_html_placeholder(title="Unknown Tags", content="No Unknown Tags Found")
525
+
526
+ return html_content, bad_entities, known_entities_in_prompt # Return list of lists for Dataframe
527
+
528
+
529
+ def build_tag_offsets_dicts(new_image_tags_with_positions):
530
+ # Structure the data for HighlightedText
531
+ tag_data = []
532
+ for tag_text, start_pos, nodetype in new_image_tags_with_positions:
533
+ # Modify the tag
534
+ modified_tag = tag_text.replace('_', ' ').replace('\\(', '(').replace('\\)', ')').strip()
535
+ artist_matrix_tag = tag_text.replace('_', ' ').replace('\\(', '\(').replace('\\)', '\)').strip()
536
+ tf_idf_matrix_tag = re.sub(r'\\([()])', r'\1', re.sub(r' ', '_', tag_text.strip().removeprefix('by ').removeprefix('by_')))
537
+ # Calculate the end position based on the original tag length
538
+ end_pos = start_pos + len(tag_text)
539
+ # Append the structured data for each tag
540
+ tag_data.append({
541
+ "original_tag": tag_text,
542
+ "start_pos": start_pos,
543
+ "end_pos": end_pos,
544
+ "modified_tag": modified_tag,
545
+ "artist_matrix_tag": artist_matrix_tag,
546
+ "tf_idf_matrix_tag": tf_idf_matrix_tag,
547
+ "node_type": nodetype
548
+ })
549
+ return tag_data
550
+
551
+
552
+ def augment_bad_entities_with_regex(text):
553
+ bad_entities = []
554
+
555
+ #comma at end
556
+ match = re.search(r',(?=\s*$)', text)
557
+ if match:
558
+ index = match.start()
559
+ bad_entities.append({"entity":"Remove Final Comma", "start":index, "end":index+1})
560
+ match = re.search(r'\([^()]*(,)\s*\)\s*$', text)
561
+ if match:
562
+ index = match.start(1)
563
+ bad_entities.append({"entity":"Remove Final Comma", "start":index, "end":index+1})
564
+ match = re.search(r'\([^()]*(,)\s*:\s*\d+(\.\d+)?\s*\)\s*$', text)
565
+ if match:
566
+ index = match.start(1)
567
+ bad_entities.append({"entity":"Remove Final Comma", "start":index, "end":index+1})
568
+
569
+ # Comma after parentheses, multiple occurrences
570
+ for match in re.finditer(r'(?<!\\)\)\s*(,)\s*[^\s]', text):
571
+ index = match.start(1)
572
+ bad_entities.append({"entity": "Move Comma Inside Parentheses", "start": index, "end": index + 1})
573
+
574
+ # Double Comma detection
575
+ for match in re.finditer(r',\s*,', text):
576
+ index = match.start()
577
+ bad_entities.append({"entity": "Double Comma", "start": index, "end": index + match.end() - match.start()})
578
+
579
+ return bad_entities
580
+
581
+ def escape_html(text):
582
+ return text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;").replace('"', "&quot;").replace("'", "&#039;")
583
+
584
+ def format_annotated_html(bad_entities, known_entities, text):
585
+ tooltip_map = {
586
+ "Unknown Tag": "This may not be a valid e621 tag. Consider removing or replacing it with tag(s) from the \"Unknown Tags\" section.",
587
+ "Duplicate": "This tag has appeared multiple times in your prompt. Consider removing the copies.",
588
+ "Remove Final Comma": "There should be no comma at the end of your prompt. Consider removing it.",
589
+ "Move Comma Inside Parentheses": "In most e621-based models, the comma following a tag functions as an &quot;attention anchor&quot;, carrying most of the tag&apos;s information. It should therefore be assigned the same weight as the rest of the tag. So instead of &quot;(lineless:1.1),&quot;, consider &quot;(lineless,:1.1)&quot; or &quot;(lineless,)&quot;",
590
+ "Double Comma": "One comma between tags is considered ample."
591
+ }
592
+ color_map = {
593
+ "Unknown Tag": ("white", "red"), # White text on red background
594
+ "Duplicate": ("black", "yellow"), # Black text on yellow background
595
+ "Remove Final Comma": ("white", "blue"), # White text on blue background
596
+ "Move Comma Inside Parentheses": ("white", "green"), # White text on green background
597
+ "Double Comma": ("white","orange")
598
+ }
599
+
600
+ # Combine and sort entities
601
+ combined_entities = bad_entities + known_entities
602
+ combined_entities = sorted(combined_entities, key=lambda x: x['start'],reverse=True)
603
+
604
+ # Generate HTML for the main text
605
+ html_text = text
606
+ for entity in combined_entities:
607
+ start = entity['start']
608
+ end = entity['end']
609
+ label = entity['entity']
610
+ if label == "Known Tag":
611
+ wiki_url = entity.get('wiki_url', '')
612
+ count = entity['count']
613
+ wiki_entry = entity.get('wiki_entry', '')
614
+ sanitized_wiki_entry = escape_html(wiki_entry) if wiki_entry else 'Unavailable'
615
+ if wiki_url: # Check if wiki_url is not empty
616
+ html_part = f'<a href="{wiki_url}" target="_blank" title="Count: {count}\tWiki: {sanitized_wiki_entry}" style="text-decoration: none; cursor: pointer; font-style: italic;">{text[start:end]}</a>'
617
+ else:
618
+ html_part = f'<span title="Count: {count}\tWiki: {sanitized_wiki_entry}" style="text-decoration: none; cursor: help; font-style: italic;">{text[start:end]}</span>'
619
+ else:
620
+ color = color_map.get(label, ("black", "white"))
621
+ html_part = f'<span style="background-color: {color[1]}; color: {color[0]};">{text[start:end]}</span>'
622
+ html_text = html_text[:start] + html_part + html_text[end:]
623
+
624
+ # Generate HTML for the color key
625
+ color_key_html = "<div style='text-align: right; margin-top: 20px;'>Key:"
626
+ used_labels = set(entity['entity'] for entity in bad_entities)
627
+ for label, colors in color_map.items():
628
+ if label in used_labels:
629
+ tooltip = tooltip_map.get(label, "")
630
+ # Adding margin-right for spacing between items
631
+ color_key_html += f" <span style='background-color: {colors[1]}; color: {colors[0]}; margin-right: 10px;' title='{tooltip}'>{label}</span>"
632
+ color_key_html += "</div>"
633
+
634
+ return f'<div style="padding: 10px; font-size: 16px;">{html_text}</div>{color_key_html}'
635
+
636
+
637
+ def find_similar_artists(original_tags_string, top_n, context_similarity_weight, allow_nsfw_tags):
638
+ try:
639
+ new_tags_string = original_tags_string.lower()
640
+ new_tags_string, removed_tags = remove_special_tags(new_tags_string)
641
+
642
+ # Parse the prompt
643
+ parsed = parser.parse(new_tags_string)
644
+ # Extract tags from the parsed tree
645
+ new_image_tags = extract_tags(parsed)
646
+ tag_data = build_tag_offsets_dicts(new_image_tags)
647
+
648
+ #Suggested tags stuff
649
+ suggested_tags_html_content = "<div class=\"scrollable-content\" style='display: inline-block; margin: 20px; text-align: center;'>"
650
+ suggested_tags_html_content += "<h1>Suggested Tags</h1>" # Heading for the table
651
+ suggested_tags = get_tfidf_reduced_similar_tags([item["tf_idf_matrix_tag"] for item in tag_data] + removed_tags, allow_nsfw_tags)
652
+
653
+ unseen_tags_data, bad_entities, known_entities = find_similar_tags(tag_data, suggested_tags, context_similarity_weight, allow_nsfw_tags)
654
+
655
+ #Bad tags stuff
656
+ bad_entities.extend(augment_bad_entities_with_regex(new_tags_string))
657
+ bad_entities.sort(key=lambda x: x['start'])
658
+ #bad_tags_illustrated_string = {"text":new_tags_string, "entities":bad_entities}
659
+ bad_tags_illustrated_html = format_annotated_html(bad_entities, known_entities, new_tags_string)
660
+
661
+ # Create a set of tags that should be filtered out
662
+ filter_tags = {entry["original_tag"].strip() for entry in tag_data}
663
+ # Use this set to filter suggested_tags
664
+ suggested_tags_filtered = OrderedDict((k, v) for k, v in suggested_tags.items() if k not in filter_tags)
665
+
666
+ # Splitting the dictionary into two based on the condition
667
+ suggested_artist_tags_filtered = OrderedDict((k, v) for k, v in suggested_tags_filtered.items() if k.startswith("by "))
668
+ suggested_non_artist_tags_filtered = OrderedDict((k, v) for k, v in suggested_tags_filtered.items() if not k.startswith("by ") and k not in special_tags)
669
+
670
+ topnsuggestions = list(islice(suggested_non_artist_tags_filtered.items(), 100))
671
+ suggested_tags_html_content += create_html_tables_for_tags("-", "Suggested Tag", topnsuggestions, find_similar_tags.tag2count, find_similar_tags.tag2idwiki)
672
+
673
+ #Artist stuff
674
+ excluded_artists = ["by conditional dnp", "by unknown artist"]
675
+ top_artists = [(key, value) for key, value in suggested_artist_tags_filtered.items() if key.lower() not in excluded_artists][:top_n]
676
+ top_artists_str = create_top_artists_table(top_artists)
677
+ dynamic_prompts_formatted_artists = "{" + "|".join([artist for artist, _ in top_artists]) + "}"
678
+
679
+ image_galleries = []
680
+ for root, dirs, files in os.walk(sample_images_directory_path):
681
+ for name in dirs:
682
+ baseline, artists = generate_artist_image_tuples([name[3:] for name, _ in top_artists], os.path.join(root, name))
683
+ image_galleries.append(baseline) # Add baseline as its own gallery item
684
+ image_galleries.append(artists) # Extend the list with artist tuples
685
+
686
+ return (unseen_tags_data, bad_tags_illustrated_html, suggested_tags_html_content, top_artists_str, dynamic_prompts_formatted_artists, *image_galleries)
687
+ except ParseError as e:
688
+ return [], "Parse Error: Check for mismatched parentheses or something", "", "", None, None
689
+
690
+
691
+ with gr.Blocks(css=css) as app:
692
+ with gr.Group():
693
+ with gr.Row():
694
+ with gr.Column(scale=3):
695
+ image_tags = gr.Textbox(label="Enter Prompt", placeholder="e.g. fox, outside, detailed background, ...")
696
+ #bad_tags_illustrated_string = gr.HighlightedText(show_legend=True, color_map={"Unknown Tag":"red","Duplicate":"yellow","Remove Final Comma":"purple","Move Comma Inside Parentheses":"green"}, label="Annotated Prompt")
697
+ bad_tags_illustrated_string = gr.HTML()
698
+ with gr.Column(scale=1):
699
+ #image_path = os.path.join("https://huggingface.co/spaces/FoodDesert/Prompt_Squirrel/resolve/main", "transparentsquirrel.png")
700
+ #gr.Image(label=" ", value=image_path, height=155, width=140)
701
+ gr.HTML('<div style="text-align: center;"><img src="https://huggingface.co/spaces/FoodDesert/Prompt_Squirrel/resolve/main/mascotimages/transparentsquirrel.png" alt="Cute Mascot" style="height: 220px; width: auto; background: transparent;"></div><br>')
702
+ #gr.HTML("<br>" * 2) # Adjust the number of line breaks ("<br>") as needed to push the button down
703
+ #image_path = os.path.join('mascotimages', "transparentsquirrel.png")
704
+ #random_image_path = os.path.join('mascotimages', random.choice([f for f in os.listdir('mascotimages') if os.path.isfile(os.path.join('mascotimages', f))]))
705
+ #with Image.open(random_image_path) as img:
706
+ # gr.Image(value=img,show_label=False, show_download_button=False, show_share_button=False, height=200)
707
+ #gr.Image(value="https://huggingface.co/spaces/FoodDesert/Prompt_Squirrel/resolve/main/mascotimages/transparentsquirrel.png",show_label=False, show_download_button=False, show_share_button=False, height=200)
708
+ #I posted the image to discord, and that's where this link came from. This is a very ugly way to do this, but I could not, no matter what I tried, get it to display an image from within the space itself. The galleries work fine for some reason, but not this.
709
+ #gr.Image(value="https://res.cloudinary.com/dnse84ol6/image/upload/v1713538125/transparentsquirrel_zhou7f.png",show_label=False, show_download_button=False, show_share_button=False, height=200)
710
+ submit_button = gr.Button(variant="primary")
711
+ with gr.Row():
712
+ with gr.Column(scale=3):
713
+ with gr.Group():
714
+ with gr.Row():
715
+ context_similarity_weight = gr.Slider(minimum=0, maximum=1, value=0.5, step=0.1, label="Context Similarity Weight")
716
+ allow_nsfw = gr.Checkbox(label="Allow NSFW Tags", value=False)
717
+ with gr.Row():
718
+ with gr.Column(scale=2):
719
+ unseen_tags = gr.HTML(label="Unknown Tags", value=create_html_placeholder(title="Unknown Tags"))
720
+ with gr.Column(scale=1):
721
+ suggested_tags = gr.HTML(label="Suggested Tags", value=create_html_placeholder(title="Suggested Tags"))
722
+ with gr.Column(scale=1):
723
+ with gr.Group():
724
+ num_artists = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Number of artists")
725
+ top_artists = gr.HTML(label="Top Artists", value=create_html_placeholder(title="Top Artists"))
726
+ dynamic_prompts = gr.Textbox(label="Dynamic Prompts Format", info="For if you're using the Automatic1111 webui (https://github.com/AUTOMATIC1111/stable-diffusion-webui) with the Dynamic Prompts extension activated (https://github.com/adieyal/sd-dynamic-prompts) and want to try them all individually.")
727
+ galleries = []
728
+ for root, dirs, files in os.walk(sample_images_directory_path):
729
+ for name in dirs:
730
+ with gr.Row():
731
+ baseline = gr.Gallery(allow_preview=False, rows=1, columns=1, height=420, scale=3)
732
+ styles = gr.Gallery(preview=False, rows=2, columns=5, height=420, scale=8)
733
+ galleries.extend([baseline, styles])
734
+
735
+ submit_button.click(
736
+ find_similar_artists,
737
+ inputs=[image_tags, num_artists, context_similarity_weight, allow_nsfw],
738
+ outputs=[unseen_tags, bad_tags_illustrated_string, suggested_tags, top_artists, dynamic_prompts] + galleries
739
+ )
740
+
741
+ gr.Markdown(faq_content)
742
+
743
+
744
+ app.queue().launch(server_name="0.0.0.0", share=True)