app.py

import gradio as gr
from sklearn.metrics.pairwise import cosine_similarity
from scipy.sparse import csr_matrix
import numpy as np
from joblib import load
import h5py
from io import BytesIO
import csv
import re
import random
import compress_fasttext
from collections import OrderedDict
from lark import Lark
from lark import Token
from lark.exceptions import ParseError
import json
import zipfile
from PIL import Image
import io
import os
    



faq_content="""
# Questions:

## What is the purpose of this tool?

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?

No

## Should I use underscores or spaces in the input tags?

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?

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.  
An example that illustrates acceptable parentheses and weight formatting is:
((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)

## Why are some valid tags marked as "unseen", and why don't some artists ever get returned?

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.
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"
in your list to bias the output toward artists with higher or lower scoring images.  Since they are not real tags, the Unseen Tags section will complain, but you can ignore that.

## Are there any other special tricks?

Yes.  If you want to more strongly bias the artist output toward a specific tag, you can just list it multiple times.  
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"
than the query "red fox, score:7".  

## Why is this space tagged "not-for-all-audience"
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.
The app will try not to display nsfw tags unless the "Allow NSFW Tags" is checked, but the filter is not perfect.

## How does the tag corrector work?

We collect the tag sets from over 4 million e621 posts, treating the tag set from each image as an individual document.
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 
(e.g. "canine" gets replaced with one of {k9,canines,mongrel,cannine,cnaine,feral_canine,anthro_canine}).
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
the vector for a tag and the vectors for its aliases are all close together (because the model has seen them in similar contexts).
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.  

To enhance the tag corrector further, we leverage conditional probabilities to refine our predictions. 
Using the same 4 million post dataset, we calculate the conditional probability of each tag given the context of other tags appearing within the same document.
This is done by creating a co-occurrence matrix from our dataset, which records how frequently each pair of tags appears together across all documents.
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.
The "similarity weight" slider controls how much weight these conditional probabilities are given vs how much weight the FastText similarity model is given when suggesting replacements for invalid tags.
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 conditional probabilities are used (although the FastText model is still used to trim the list of candidates).

## How is the artist list calculated?

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. 
Similarly, when you input a set of tags, the system creates a pseudo-document for your query out of all the tags. 
It then uses a technique called cosine similarity to compare your tags against each artist's collection, essentially finding which artist's tags are most "similar" to yours.
This method helps identify artists whose work is closely aligned with the themes or elements you're interested in.
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. 
You can read more about TF-IDF on its [Wikipedia page](https://en.wikipedia.org/wiki/Tf%E2%80%93idf).  

## How do the sample images work?

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".
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.
The first image returned is a baseline, generated with the same prompt, but with no artist name.
You should compare all the images to the first to see how the artist names affect the output.

## Why is the layout so messed up?

I do not know how to resolve an issue with the Gallery.  I shuffled all the components around so that it would be at the bottom and not interfere with anything else.
"""


nsfw_threshold = 0.95  # Assuming the threshold value is defined here

grammar=r"""
!start: (prompt | /[][():]/+)*
prompt: (emphasized | plain | comma | WHITESPACE)*
!emphasized: "(" prompt ")"
        | "(" prompt ":" [WHITESPACE] NUMBER [WHITESPACE] ")"
comma: ","
WHITESPACE: /\s+/
plain: /([^,\\\[\]():|]|\\.)+/
%import common.SIGNED_NUMBER -> NUMBER
"""
# Initialize the parser
parser = Lark(grammar, start='start')

# Function to extract tags
def extract_tags(tree):
    tags = []
    def _traverse(node):
        if isinstance(node, Token) and node.type == '__ANON_1':
            tags.append(node.value.strip())
        elif not isinstance(node, Token):
            for child in node.children:
                _traverse(child)

    _traverse(tree)
    return tags


special_tags = ["score:0", "score:1", "score:2", "score:3", "score:4", "score:5", "score:6", "score:7", "score:8", "score:9"]
def remove_special_tags(original_string):
    tags = [tag.strip() for tag in original_string.split(",")]
    remaining_tags = [tag for tag in tags if tag not in special_tags]
    removed_tags = [tag for tag in tags if tag in special_tags]
    return ", ".join(remaining_tags), removed_tags
    
    
# Load the model and data once at startup
with h5py.File('complete_artist_data.hdf5', 'r') as f:
    # Deserialize the vectorizer
    vectorizer_bytes = f['vectorizer'][()].tobytes()
    # Use io.BytesIO to convert bytes back to a file-like object for joblib to load
    vectorizer_buffer = BytesIO(vectorizer_bytes)
    vectorizer = load(vectorizer_buffer)
    
    # Load X_artist
    X_artist = f['X_artist'][:]
    # Load artist names and decode to strings
    artist_names = [name.decode() for name in f['artist_names'][:]]


with h5py.File('conditional_tag_probabilities_matrix.h5', 'r') as f:
    # Reconstruct the sparse co-occurrence matrix
    conditional_co_occurrence_matrix = csr_matrix(
        (f['co_occurrence_data'][:], f['co_occurrence_indices'][:], f['co_occurrence_indptr'][:]),
        shape=f['co_occurrence_shape'][:]
    )

    # Reconstruct the vocabulary
    conditional_words = f['vocabulary_words'][:]
    conditional_indices = f['vocabulary_indices'][:]
    conditional_vocabulary = {key.decode('utf-8'): value for key, value in zip(conditional_words, conditional_indices)}

    # Load the document count
    conditional_doc_count = f['doc_count'][()]
    conditional_smoothing = 100. / conditional_doc_count
    
    
nsfw_tags = set()  # Initialize an empty set to store words meeting the threshold
# Open and read the CSV file
with open("word_rating_probabilities.csv", 'r', newline='', encoding='utf-8') as csvfile:
    reader = csv.reader(csvfile)
    next(reader, None)  # Skip the header row
    for row in reader:
        word = row[0]  # The word is in the first column
        probability_sum = float(row[1])  # The sum of probabilities is in the second column, convert to float for comparison
        # Check if the probability sum meets the threshold and add the word to the set if it does
        if probability_sum >= nsfw_threshold:
            nsfw_tags.add(word)
    
    
soyjak_directory_path = 'artistsoyjaks'
soyjak_json_file_path = os.path.join(soyjak_directory_path, 'artistsoyjaks.json')
with open(soyjak_json_file_path, 'r') as json_file:
    soyjak_artist_to_file_map = json.load(json_file)

    
def generate_artist_image_tuples(top_artists, image_directory=soyjak_directory_path):
    artist_image_tuples = []
    for artist in top_artists:
        filename = soyjak_artist_to_file_map.get(artist)
        if filename:
            image_path = os.path.join(image_directory, filename)
            if os.path.exists(image_path):
                artist_image_tuples.append((image_path, artist if artist else "No Artist"))

    return artist_image_tuples
    
    
def clean_tag(tag):
    return ''.join(char for char in tag if ord(char) < 128)
    
    
#Normally returns tag to aliases, but when reverse=True, returns alias to tags
def build_aliases_dict(filename, reverse=False):   
    aliases_dict = {}
    with open(filename, 'r', newline='', encoding='utf-8') as csvfile:
        reader = csv.reader(csvfile)
        for row in reader:
            tag = clean_tag(row[0])
            alias_list = [] if row[3] == "null" else [clean_tag(alias) for alias in row[3].split(',')]
            if reverse:
                for alias in alias_list:
                    aliases_dict.setdefault(alias, []).append(tag)
            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.
#Note the intention is that sum_i(P(word_i|word_j)) =(approx) # of words in a document rather than 1.
def conditional_probability(word_i, word_j, co_occurrence_matrix, vocabulary, doc_count, smoothing_value=0.01):
    word_i_index = vocabulary.get(word_i)
    word_j_index = vocabulary.get(word_j)
    
    if word_i_index is not None and word_j_index is not None:
        # Directly access the sparse matrix elements
        word_j_count = co_occurrence_matrix[word_j_index, word_j_index]
        smoothed_word_j_count =  word_j_count + (smoothing_value * doc_count)
        
        word_i_count = co_occurrence_matrix[word_i_index, word_i_index]
        
        co_occurrence_count = co_occurrence_matrix[word_i_index, word_j_index]
        smoothed_co_occurrence_count = co_occurrence_count + (smoothing_value * word_i_count) 
        
        # Calculate the conditional probability with smoothing
        conditional_prob = smoothed_co_occurrence_count / smoothed_word_j_count
        
        return conditional_prob
    elif word_i_index is None:
        return 0
    else:
        return None


#geometric_mean_given_words(target_word, context_words, conditional_co_occurrence_matrix, conditioanl_vocabulary, conditional_doc_count, smoothing_value=conditional_smoothing):
def geometric_mean_given_words(target_word, context_words, co_occurrence_matrix, vocabulary, doc_count, smoothing_value=0.01):
    probabilities = []
    
    # Collect the conditional probabilities of the target word given each context word, ignoring None values
    for context_word in context_words:
        prob = conditional_probability(target_word, context_word, co_occurrence_matrix, vocabulary, doc_count, smoothing_value)
        if prob is not None:
            probabilities.append(prob)
    
    # Compute the geometric mean of the probabilities, avoiding division by zero
    if probabilities:  # Check if the list is not empty
        geometric_mean = np.prod(probabilities) ** (1.0 / len(probabilities))
    else:
        geometric_mean = 0.5  # Or assign some default value if all probabilities are None
    
    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):
    #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')
    tag_aliases_file = 'fluffyrock_3m.csv'
    if not hasattr(find_similar_tags, "tag2aliases"):
        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 tables
    html_content = ""
    for tag in test_tags:
        if tag in special_tags:
            continue
            
        modified_tag_for_search = tag.replace(' ','_')
        similar_words = find_similar_tags.fasttext_small_model.most_similar(modified_tag_for_search, topn = 100)
        result, seen = [], set(transformed_tags)
        
        if modified_tag_for_search in find_similar_tags.tag2aliases:
            if tag in find_similar_tags.tag2aliases and "_" in tag:   #Implicitly tell the user that they should get rid of the underscore
                result.append(modified_tag_for_search.replace('_',' '), 1)
                seen.add(tag)
            else:   #The user correctly did not put underscores in their tag
                continue
        else:
            for item in similar_words:
                similar_word, similarity = item
                if similar_word not in seen:
                    if similar_word in find_similar_tags.tag2aliases:
                        result.append((similar_word.replace('_', ' '), round(similarity, 3)))
                        seen.add(similar_word)
                    else:
                        for similar_tag in find_similar_tags.alias2tags.get(similar_word, []):
                            if similar_tag not in seen:
                                result.append((similar_tag.replace('_', ' '), round(similarity, 3)))
                                seen.add(similar_tag)

        #Remove NSFW tags if appropriate.
        if not allow_nsfw_tags:
            result = [(word, score) for word, score in result if word.replace(' ','_') not in nsfw_tags]

        #Adjust score based on context
        for i in range(len(result)):
            word, score = result[i]  # Unpack the tuple
            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)

        result = sorted(result, key=lambda x: x[1], reverse=True)[:10]
        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 html_content  # Return list of lists for Dataframe

def find_similar_artists(new_tags_string, top_n, similarity_weight, allow_nsfw_tags):
    try:
        new_tags_string = new_tags_string.lower()
        new_tags_string, removed_tags = remove_special_tags(new_tags_string)
        
        # Parse the prompt
        parsed = parser.parse(new_tags_string)
        # Extract tags from the parsed tree
        new_image_tags = extract_tags(parsed)
        new_image_tags = [tag.replace('_', ' ').replace('\\(', '(').replace('\\)', ')').strip() for tag in new_image_tags]

        ###unseen_tags = list(set(OrderedDict.fromkeys(new_image_tags)) - set(vectorizer.vocabulary_.keys()))   #We may want this line again later.  These are the tags that were not used to calculate the artists list.
        unseen_tags_data = find_similar_tags(new_image_tags, similarity_weight, allow_nsfw_tags)

        X_new_image = vectorizer.transform([','.join(new_image_tags + removed_tags)])
        similarities = cosine_similarity(X_new_image, X_artist)[0]

        top_artist_indices = np.argsort(similarities)[-top_n:][::-1]
        top_artists = [(artist_names[i], similarities[i]) for i in top_artist_indices]

        top_artists_str = "\n".join([f"{rank+1}. {artist[3:]} ({score:.4f})" for rank, (artist, score) in enumerate(top_artists)])
        dynamic_prompts_formatted_artists = "{" + "|".join([artist for artist, _ in top_artists]) + "}"

        image_gallery = generate_artist_image_tuples([''] + [name[3:] for name, _ in top_artists])

        return unseen_tags_data, top_artists_str, dynamic_prompts_formatted_artists, image_gallery
    except ParseError as e:
        return [], "Parse Error: Check for mismatched parentheses or something", "", None


with gr.Blocks() as app:
    with gr.Group():
        with gr.Row():
            image_tags = gr.Textbox(label="Enter image tags", placeholder="e.g. fox, outside, detailed background, ...")
        with gr.Row():
            submit_button = gr.Button("Submit")
            similarity_weight = gr.Slider(minimum=0, maximum=1, value=0.5, step=0.1, label="Similarity weight")
            num_artists = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Number of artists")
            allow_nsfw = gr.Checkbox(label="Allow NSFW Tags", value=False)
    with gr.Row():
        with gr.Column():
            top_artists = gr.Textbox(label="Top Artists", info="These are the artists most strongly associated with your tags. The number in parentheses is a similarity score between 0 and 1, with higher numbers indicating greater similarity.")
        with gr.Column():
            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.")
    with gr.Row():
        unseen_tags = gr.HTML(label="Unseen Tags")
    with gr.Row():
        styles = gr.Gallery(label="Styles", allow_preview=True, preview=True, container=True, rows=1, columns=11)
    
    
    submit_button.click(
        find_similar_artists, 
        inputs=[image_tags, num_artists, similarity_weight, allow_nsfw], 
        outputs=[unseen_tags, top_artists, dynamic_prompts, styles]
    )
    
    gr.Markdown(faq_content)
    

app.launch()