Upload 2 files

.gitattributes

@@ -35,3 +35,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 artist_data.hdf5 filter=lfs diff=lfs merge=lfs -text
 complete_artist_data.hdf5 filter=lfs diff=lfs merge=lfs -text
+wiki_pages-2023-08-08.csv filter=lfs diff=lfs merge=lfs -text

app.py

@@ -22,10 +22,14 @@ faq_content="""
 
 ## What is the purpose of this tool?
 
-When you enter a txt2img prompt and press the "submit" button, the Tagset Completer parses your prompt and checks that all your tags are valid e621 tags.  
-If it finds any that are not, it recommends some valid e621 tags you can use to replace them in the "Unseen Tags" table.  
-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, 
-in case you want to look them up to get more ideas.  This is useful to align your prompt with the expected input to an e621-trained model.
+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.
+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.
+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.
+
+When you enter a txt2img prompt and press the "submit" button, the Tagset Completer parses your prompt and checks that all your tags are valid e621 tags.
+If it finds any that are not, it recommends some valid e621 tags you can use to replace them in the "Unseen Tags" table.
+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.
+This is useful to align your prompt with the expected input to an e621-trained model.
 
 ## Does input order matter?
 
@@ -33,7 +37,7 @@ No
 
 ## Should I use underscores or spaces in the input tags?
 
-Spaces are preferred, but it will still work if you use underscores.  The Unseen Tags table will just complain at you.
+As a rule, e621-trained models replace underscores in tags with spaces, so spaces are preferred.
 
 ## Can I use parentheses or weights as in the Stable Diffusion Automatic1111 WebUI?
 
@@ -46,6 +50,10 @@ An example that illustrates acceptable parentheses and weight formatting is:
 Some data is excluded from consideration if it did not occur frequently enough in the sample from which the application makes its calculations.
 If an artist or tag is too infrequent, we might not think we have enough data to make predictions about it.
 
+## Why do some suggested tags not have summaries or wiki links?
+
+Both of these features are extracted from the tag wiki pages, but some valid e621 tags do not have wiki pages.
+
 ## Are there any special tags?
 
 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.
@@ -188,6 +196,63 @@ def build_aliases_dict(filename, reverse=False):
             else:
                 aliases_dict[tag] = alias_list
     return aliases_dict
+    
+
+def build_tag_count_dict(filename):
+    with open(filename, 'r', newline='', encoding='utf-8') as csvfile:
+        reader = csv.reader(csvfile)
+        result_dict = {}
+        for row in reader:
+            key = row[0]
+            value = int(row[2]) if row[2].isdigit() else None
+            if value is not None:
+                result_dict[key] = value
+    return result_dict
+
+import csv
+
+
+def build_tag_id_wiki_dict(filename='wiki_pages-2023-08-08.csv'):
+    """
+    Reads a CSV file and returns a dictionary mapping tag names to tuples of
+    (number, most relevant line from the wiki entry). Rows with a non-integer in the first column are ignored.
+    The most relevant line is the first line that does not start with "thumb" and is not blank.
+
+    Parameters:
+    - filename: The path to the CSV file.
+
+    Returns:
+    - A dictionary where each key is a tag name and each value is a tuple (number, most relevant wiki entry line).
+    """
+    tag_data = {}
+    with open(filename, 'r', encoding='utf-8') as csvfile:
+        reader = csv.reader(csvfile)
+
+        # Skip the header row
+        next(reader)
+
+        for row in reader:
+            try:
+                # Attempt to convert the first column to an integer
+                number = int(row[0])
+            except ValueError:
+                # If conversion fails, skip this row
+                continue
+
+            tag = row[3]
+            wiki_entry_full = row[4]
+
+            # Process the wiki_entry to find the most relevant line
+            relevant_line = ''
+            for line in wiki_entry_full.split('\n'):
+                if line.strip() and not line.startswith("thumb"):
+                    relevant_line = line
+                    break
+
+            # Map the tag to a tuple of (number, relevant_line)
+            tag_data[tag] = (number, relevant_line)
+
+    return tag_data
 
 
 #Imagine we are adding smoothing_value to the number of times word_j occurs in each document for smoothing.
@@ -234,9 +299,32 @@ def geometric_mean_given_words(target_word, context_words, co_occurrence_matrix,
     
     return geometric_mean
 
+    
+def create_html_tables_for_tags(tag, result, tag2count, tag2idwiki):
+    # Wrap the tag part in a <span> with styles for bold and larger font
+    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;'>Unknown Tag: <span style='font-weight: bold; font-size: 20px;'>{tag}</span></th></tr></thead><tbody><tr style='border-bottom: 1px solid #000;'><th>Corrected Tag</th><th>Similarity</th><th>Count</th></tr>"
+    # Loop through the results and add table rows for each
+    for word, sim in result:
+        word_with_underscores = word.replace(' ', '_')
+        count = tag2count.get(word_with_underscores, 0)  # Get the count if available, otherwise default to 0
+        tag_id, wiki_entry = tag2idwiki.get(word_with_underscores, (None, ''))
+        # Check if tag_id and wiki_entry are valid
+        if tag_id is not None and wiki_entry:
+            # Construct the URL for the tag's wiki page
+            wiki_url = f"https://e621.net/wiki_pages/{tag_id}"
+            # Make the tag a hyperlink with a tooltip
+            tag_element = f"<a href='{wiki_url}' target='_blank' title='{wiki_entry}'>{word}</a>"
+        else:
+            # Display the word without any hyperlink or tooltip
+            tag_element = word
+        # Include the tag element in the table row
+        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>"
+
+    html_str += "</tbody></table></div>"
+    return html_str
 
-def find_similar_tags(test_tags, similarity_weight, allow_nsfw_tags):
 
+def find_similar_tags(test_tags, similarity_weight, allow_nsfw_tags):
     #Initialize stuff
     if not hasattr(find_similar_tags, "fasttext_small_model"):
         find_similar_tags.fasttext_small_model = compress_fasttext.models.CompressedFastTextKeyedVectors.load('e621FastTextModel010Replacement_small.bin')
@@ -245,11 +333,15 @@ def find_similar_tags(test_tags, similarity_weight, allow_nsfw_tags):
         find_similar_tags.tag2aliases = build_aliases_dict(tag_aliases_file)
     if not hasattr(find_similar_tags, "alias2tags"):
         find_similar_tags.alias2tags = build_aliases_dict(tag_aliases_file, reverse=True)
+    if not hasattr(find_similar_tags, "tag2count"):
+        find_similar_tags.tag2count = build_tag_count_dict(tag_aliases_file)
+    if not hasattr(find_similar_tags, "tag2idwiki"):    
+        find_similar_tags.tag2idwiki = build_tag_id_wiki_dict()
     
     transformed_tags = [tag.replace(' ', '_') for tag in test_tags]
 
-    # Find similar tags and prepare data for dataframe.
-    results_data = []
+    # Find similar tags and prepare data for tables
+    html_content = ""
     for tag in test_tags:
         if tag in special_tags:
             continue
@@ -287,22 +379,15 @@ def find_similar_tags(test_tags, similarity_weight, allow_nsfw_tags):
             geometric_mean = geometric_mean_given_words(word.replace(' ','_'), [context_tag for context_tag in transformed_tags if context_tag != word and context_tag != tag], conditional_co_occurrence_matrix, conditional_vocabulary, conditional_doc_count, smoothing_value=conditional_smoothing)
             adjusted_score = (similarity_weight * geometric_mean) + ((1-similarity_weight)*score)  # Apply the adjustment function
             result[i] = (word, adjusted_score)  # Update the tuple with the adjusted score
+            #print(word, score, geometric_mean, adjusted_score)
 
-        # Append tag and formatted similar tags to results_data
         result = sorted(result, key=lambda x: x[1], reverse=True)[:10]
-        first_entry_for_tag = True
-        for word, sim in result:
-            if first_entry_for_tag:
-                results_data.append([tag, word, sim])
-                first_entry_for_tag = False
-            else:
-                results_data.append(["", word, sim])
-        results_data.append(["", "", ""])  # Adds a blank line after each group of tags
-
-    if not results_data:
-        results_data.append(["No Unknown Tags Found", "", ""])
+        html_content += create_html_tables_for_tags(tag, result, find_similar_tags.tag2count, find_similar_tags.tag2idwiki)
+    # If no tags were processed, add a message
+    if not html_content:
+        html_content = "<p>No Unknown Tags Found</p>"
 
-    return results_data  # Return list of lists for Dataframe
+    return html_content  # Return list of lists for Dataframe
 
 def find_similar_artists(new_tags_string, top_n, similarity_weight, allow_nsfw_tags):
     try:
@@ -341,7 +426,7 @@ iface = gr.Interface(
         gr.Checkbox(label="Allow NSFW Tags", value=False)
     ],
     outputs=[
-        gr.Dataframe(label="Unseen Tags", headers=["Tag", "Similar Tags", "Similarity"]),
+        gr.HTML(label="Unseen Tags"),
         gr.Textbox(label="Top Artists", info="These are the artists most strongly associated with your tags.  The number in parenthes is a similarity score between 0 and 1, with higher numbers indicating greater similarity."),
         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.") 
     ],

wiki_pages-2023-08-08.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d453c0cc8ae09c548e554ceb77b1c1578c277eb2c5a6278a85f89c73566a7b27
+size 30986436