Spaces:
Running
Running
| import gradio as gr | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| from scipy.sparse import csr_matrix | |
| import numpy as np | |
| import joblib | |
| 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, Tree, Token | |
| from lark.exceptions import ParseError | |
| import json | |
| import zipfile | |
| from PIL import Image | |
| import io | |
| import os | |
| import glob | |
| import itertools | |
| from itertools import islice | |
| from pathlib import Path | |
| import logging | |
| import hnswlib | |
| import pathlib | |
| from collections import Counter | |
| # Set up logging | |
| # Minimal prod logging: warnings+ to stderr, no file by default | |
| import os, logging | |
| LOG_LEVEL = os.environ.get("PSQ_LOG_LEVEL", "WARNING").upper() | |
| logging.basicConfig( | |
| level=getattr(logging, LOG_LEVEL, logging.WARNING), | |
| format="%(asctime)s %(levelname)s:%(message)s", | |
| handlers=[logging.StreamHandler()] # no file -> avoids huge logs on Spaces | |
| ) | |
| # Quiet down common noisy libs (optional) | |
| for _name in ("gensim", "gradio", "hnswlib", "httpx", "uvicorn"): | |
| logging.getLogger(_name).setLevel(logging.ERROR) | |
| # Turn off Gradio analytics phone-home to avoid those background thread errors (optional) | |
| os.environ["GRADIO_ANALYTICS_ENABLED"] = "0" | |
| MASCOT_DIR = Path(__file__).parent / "mascotimages" | |
| MASCOT_FILE = MASCOT_DIR / "transparentsquirrel.png" | |
| faq_content=""" | |
| # Quick Start | |
| Write your prompt as a simple comma‑separated list of things you want to see in your image, then press the Run button. Prompt Squirrel will: | |
| * highlight any unknown or misspelled tags or syntax errors, | |
| * suggest corrected tags, | |
| * recommend additional tags based on context, and | |
| * list artists who produce topically similar content | |
| You do not need to structure the prompt in any special way; just describe what you want in short phrases separated by commas. | |
| # System Overview | |
| Prompt Squirrel uses NLP and vector‑space methods to map a free‑form prompt to the structured tag vocabulary expected by tag‑based Stable Diffusion models. Internally, we use a grammar parser, FastText embeddings, TF‑IDF and SVD for context scoring, and an approximate‑nearest‑neighbor index for artist and suggested tag retrieval. Our goal is to help users write prompts that align with the tag distributions the model was trained on. | |
| See the Technical Details heading below for more information about how these all are used. | |
| # Prompting Guidance and Common Questions | |
| ## What text to image models does this tool work for? | |
| This instance of Prompt Squirrel works for tag-based Stable Diffusion models fine-tuned on the popular e621 dataset. The tags it returns and especially the artist names will only be recognized by models in this category, which includes popular models such as Fluffyrock and NoobAI. | |
| ## Does input order matter? | |
| No. | |
| ## Should I use underscores or spaces in the input tags? | |
| As a rule, models trained on the dataset replace underscores 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. These are ignored in the underlying calculations but allowed so that prompts can be copied between tools with minimal editing. An acceptable example 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 'unknown', and why don't some artists ever get returned? | |
| Some tags or artists are too infrequent in the dataset sample for us to make reliable predictions. Prompt Squirrel merges data from several sources, which may differ slightly in tag names or counts. Low‑frequency items or inconsistent entries may therefore not appear in results. | |
| ## Why do some suggested tags not have summaries or wiki links, and why do some look truncated? | |
| Summaries and wiki links are extracted from dataset wiki pages. Some tags do not have pages, and summaries are heuristically extracted from the page beginnings, which can introduce small errors. | |
| ## Are there any special tags? | |
| Yes. We normalized favorite counts to a range of 0–9. You may include: 'score:0' through 'score:9' These bias the output toward suggestions associated with images with higher or lower scores. | |
| ## Are there any other special tricks? | |
| Yes. Repeating a tag gives it more weight in the artist‑similarity calculation. For example: 'red fox, red fox, red fox, score:7' will bias more strongly toward artists and suggested tags associated with 'red fox' than: 'red fox, score:7'. | |
| ## Why is this space tagged 'not‑for‑all‑audience'? | |
| The dataset used by many tag‑based models contains both general‑audience and adult material. To avoid surprising users, mature tags are hidden unless the user explicitly enables them. This tool processes only text and metadata; no images from the dataset are displayed. | |
| # Technical Details | |
| ## How is the artist list calculated? | |
| Each artist is represented as a pseudo‑document containing the bag of all tags from their images. Your prompt is treated as another pseudo‑document. We compute similarity between the recognized tags in your prompt and each artist using TF‑IDF and truncated SVD, then retrieve the nearest artists using an approximate‑nearest‑neighbor index. | |
| ## How do the Suggested Tags work? | |
| Each tag is represented as a pseudo-document containing the bag of all tags it co-occurs with in the dataset. We then employ exactly the same method on them as we did with artists to suggest tags similar to your prompt. | |
| ## How does the tag corrector work? | |
| We treat each image as a document containing the set of its tags and randomly replace about 10% of the tags with aliases from the dataset's alias lists. We then train a FastText model on these documents so that tags and their aliases map to nearby vectors. This makes the system robust to spelling variations and rephrasings. | |
| To incorporate the prompt's context, we again treat tags as pseudo-documents containing the bag of all tags they co-occur with, then compute TF‑IDF scores for the top candidate tags selected from the FastText Model, combining this similarity score with the FastText similarity. The Context Similarity Weight slider controls how much influence the TF‑IDF context score has relative to the FastText embedding similarity. | |
| ## How do the sample images work? | |
| In each gallery row, we choose an illustrative prompt and generate one image for each artist using the popular model Fluffyrock Unleashed, which was trained on this dataset. The 'No Artist' image serves as a baseline, using the same prompt without an artist name. Each subsequent row repeats this process with a different prompt. The first prompt was chosen to illustrate foreground style, the second to illustrate background style, and the third to illustrate character design. See SamplePrompts.csv in the Files section for the list of prompts used. | |
| """ | |
| TOOLTIP_NOTE_HTML = '<div class="hover-hint">Hover over underlined items for more info.</div>' | |
| HOVER_HINT_CSS = """ | |
| /* Solid, visible underline for tagged items */ | |
| .gradio-container .hover-underline{ | |
| text-decoration-line: underline !important; | |
| text-decoration-thickness: 2px; | |
| text-underline-offset: 2px; | |
| } | |
| /* Small, subtle hint text */ | |
| .hover-hint{ | |
| font-size: 12px; | |
| opacity: .85; | |
| line-height: 1.2; | |
| } | |
| /* Wrapper to position the hint in the bottom-right of the annotated box */ | |
| .annotated-wrap{ position: relative; } | |
| .annotated-wrap .hover-hint{ | |
| position: absolute; | |
| right: 6px; | |
| bottom: 6px; | |
| text-align: right; | |
| } | |
| /* Extra padding around commas */ | |
| .badpad{ | |
| display:inline-block; /* make padding take effect */ | |
| padding: 0 0.60em; /* ~two extra spaces total on each side of ',' -> ~3-ch width overall */ | |
| border-radius: 3px; /* a tiny rounding looks nicer on a blocky background */ | |
| } | |
| """ | |
| try: | |
| from gradio_client import utils as _gc_utils | |
| _orig_get_type = _gc_utils.get_type | |
| _orig_j2p = _gc_utils._json_schema_to_python_type | |
| _orig_pub = _gc_utils.json_schema_to_python_type | |
| def _get_type_safe(schema): | |
| # Sometimes schema is a bare True/False (JSON Schema boolean form) | |
| if not isinstance(schema, dict): | |
| return "any" | |
| return _orig_get_type(schema) | |
| def _j2p_safe(schema, defs=None): | |
| # Accept non-dict schemas (True/False/None) and treat as "any" | |
| if not isinstance(schema, dict): | |
| return "any" | |
| return _orig_j2p(schema, defs or schema.get("$defs")) | |
| def _pub_safe(schema): | |
| # Public wrapper used by Gradio; keep it resilient too | |
| if not isinstance(schema, dict): | |
| return "any" | |
| return _j2p_safe(schema, schema.get("$defs")) | |
| _gc_utils.get_type = _get_type_safe | |
| _gc_utils._json_schema_to_python_type = _j2p_safe | |
| _gc_utils.json_schema_to_python_type = _pub_safe | |
| except Exception as e: | |
| print("gradio_client hotfix not applied:", e) | |
| # ------------------------------------------------------------------------------- | |
| nsfw_threshold = 0.95 # Assuming the threshold value is defined here | |
| css = HOVER_HINT_CSS + """ | |
| .scrollable-content{ | |
| max-height: 420px; | |
| overflow-y: scroll; /* always show scrollbar */ | |
| overflow-x: hidden; | |
| padding-right: 8px; | |
| padding-bottom: 14px; /* <— add this */ | |
| scrollbar-gutter: stable; /* prevent layout shift as it fills */ | |
| /* Firefox */ | |
| scrollbar-width: auto; | |
| scrollbar-color: rgba(180,180,180,.9) rgba(0,0,0,.15); | |
| } | |
| /* WebKit/Chromium (Chrome/Edge/Safari) */ | |
| .scrollable-content::-webkit-scrollbar{ width: 10px; } | |
| .scrollable-content::-webkit-scrollbar-thumb{ background: rgba(180,180,180,.9); border-radius: 8px; } | |
| .scrollable-content::-webkit-scrollbar-track{ background: rgba(0,0,0,.15); } | |
| /* (Optional) make both scroll panes taller so they fill more of the column */ | |
| .pane-left .scrollable-content, | |
| .pane-right .scrollable-content { | |
| max-height: 610px; /* was 420px; tweak to taste */ | |
| } | |
| """ | |
| #Parser | |
| 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') | |
| # ---------- Two HNSW indexes: artists and non-artist tags ---------- | |
| _HNSW_ART = None | |
| _HNSW_TAG = None | |
| _HNSW_DIM = None | |
| _HNSW_N_ART = None | |
| _HNSW_N_TAG = None | |
| _HNSW_ART_PATH = pathlib.Path("tfidf_hnsw_artists.bin") | |
| _HNSW_TAG_PATH = pathlib.Path("tfidf_hnsw_tags.bin") | |
| def _l2_normalize_rows(mat: np.ndarray) -> np.ndarray: | |
| mat = np.asarray(mat, dtype=np.float32) | |
| norms = np.linalg.norm(mat, axis=1, keepdims=True) | |
| norms[norms == 0.0] = 1.0 | |
| return mat / norms | |
| def _ensure_dual_hnsw_indexes(): | |
| """ | |
| Build/load two HNSW indexes over the SVD-reduced TF-IDF matrix: | |
| • _HNSW_ART — rows whose tag (with optional 'by_' stripped) is in the artist_set | |
| • _HNSW_TAG — only rows that are NOT artist tags | |
| Index item IDs are the ORIGINAL row indices in reduced_matrix. | |
| """ | |
| global _HNSW_ART, _HNSW_TAG, _HNSW_DIM, _HNSW_N_ART, _HNSW_N_TAG | |
| if _HNSW_ART is not None and _HNSW_TAG is not None: | |
| return | |
| reduced_matrix = tf_idf_components['reduced_matrix'] # (N, D) | |
| row_to_tag = tf_idf_components['row_to_tag'] # {row:int -> "tag_with_underscores"} | |
| rm = _l2_normalize_rows(reduced_matrix).astype(np.float32) | |
| n_items, dim = rm.shape | |
| # Partition rows | |
| artist_rows = [] | |
| tag_rows = [] | |
| for i in range(n_items): | |
| tag = row_to_tag.get(i, "") | |
| # Strip leading "by_" if present in the TF-IDF vocabulary, but don't rely on it. | |
| base = tag[3:] if tag.startswith("by_") else tag | |
| # Some corpora contain buckets you don't want shown as artists: | |
| if tag in {"by_unknown_artist", "by_conditional_dnp"}: | |
| tag_rows.append(i) | |
| continue | |
| if is_artist(base): | |
| artist_rows.append(i) | |
| else: | |
| tag_rows.append(i) | |
| logging.debug(f"HNSW partition: artists={len(artist_rows)} non_artists={len(tag_rows)}") | |
| # Helper: build or load an index for a subset of rows | |
| def _build_or_load(path: pathlib.Path, rows: list[int]) -> hnswlib.Index: | |
| idx = hnswlib.Index(space='cosine', dim=dim) | |
| need_build = True | |
| if path.exists(): | |
| try: | |
| idx.load_index(str(path), max_elements=max(1, len(rows))) | |
| # Rebuild if the saved index count doesn’t match our rows | |
| if getattr(idx, "get_current_count", None) and idx.get_current_count() == len(rows) and len(rows) > 0: | |
| need_build = False | |
| else: | |
| logging.debug(f"Rebuilding {path.name}: saved_count!=rows_len ({idx.get_current_count()} vs {len(rows)})") | |
| except Exception as e: | |
| logging.debug(f"Reload {path.name} failed, rebuilding: {e}") | |
| if need_build: | |
| try: | |
| if path.exists(): | |
| path.unlink() | |
| except Exception: | |
| pass | |
| idx.init_index(max_elements=max(1, len(rows)), ef_construction=200, M=16) | |
| if rows: | |
| idx.add_items(rm[rows], ids=np.asarray(rows, dtype=np.int32)) | |
| idx.save_index(str(path)) | |
| idx.set_ef(200) | |
| return idx | |
| _HNSW_ART = _build_or_load(_HNSW_ART_PATH, artist_rows) | |
| _HNSW_TAG = _build_or_load(_HNSW_TAG_PATH, tag_rows) | |
| _HNSW_DIM = dim | |
| _HNSW_N_ART = len(artist_rows) | |
| _HNSW_N_TAG = len(tag_rows) | |
| def _hnsw_query(idx: hnswlib.Index, vec: np.ndarray, k: int): | |
| """ | |
| Query a given HNSW index with a (1, D) or (D,) vector in SVD space. | |
| Returns (indices, sims) with cosine similarity scores. | |
| """ | |
| _ensure_dual_hnsw_indexes() | |
| q = np.asarray(vec, dtype=np.float32).reshape(-1) | |
| q_norm = np.linalg.norm(q) | |
| if q_norm > 0: | |
| q = q / q_norm | |
| labels, dists = idx.knn_query(q, k=k) | |
| inds = labels[0] | |
| sims = 1.0 - dists[0] # cosine distance -> similarity | |
| return inds, sims | |
| def _ann_tags_topk(vec: np.ndarray, k: int): | |
| _ensure_dual_hnsw_indexes() | |
| k = min(k, _HNSW_N_TAG if _HNSW_N_TAG else 0) | |
| return _hnsw_query(_HNSW_TAG, vec, k) if k else (np.array([], dtype=int), np.array([], dtype=float)) | |
| def _ann_artists_topk(vec: np.ndarray, k: int): | |
| _ensure_dual_hnsw_indexes() | |
| k = min(k, _HNSW_N_ART if _HNSW_N_ART else 0) | |
| return _hnsw_query(_HNSW_ART, vec, k) if k else (np.array([], dtype=int), np.array([], dtype=float)) | |
| # ------------------------------------------------------------------ | |
| def _norm_tag_for_lookup(s: str) -> str: | |
| # convert "name with spaces" -> "name_with_spaces" and unescape parens | |
| return s.replace(' ', '_').replace('\\(', '(').replace('\\)', ')') | |
| # Function to extract tags | |
| def extract_tags(tree): | |
| tags_with_positions = [] | |
| def _traverse(node): | |
| if isinstance(node, Token) and node.type == '__ANON_1': | |
| tag_position = node.start_pos | |
| tag_text = node.value | |
| tags_with_positions.append((tag_text, tag_position, "tag")) | |
| elif not isinstance(node, Token): | |
| for child in node.children: | |
| _traverse(child) | |
| _traverse(tree) | |
| return tags_with_positions | |
| 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"] | |
| 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 | |
| #Model specific tags | |
| MODEL_SPECIFIC_TAGS = { | |
| "masterpiece", | |
| "best quality", | |
| "good quality", | |
| "normal quality", | |
| "newest", | |
| "absurdres", | |
| "highres", | |
| "safe", | |
| "worst quality", | |
| "early", | |
| "low quality", | |
| "lowres", | |
| "explict content", | |
| "very awa", | |
| "worst aesthetic", | |
| "score_9", | |
| "score_8_up", | |
| "score_7_up", | |
| "score_6_up", | |
| "score_5_up", | |
| "score_4_up", | |
| "source_pony", | |
| "source_furry", | |
| "source_cartoon", | |
| "source_anime", | |
| "rating_safe", | |
| "rating_questionable", | |
| "rating_explicit" | |
| } | |
| # Define a function to load all necessary components | |
| def load_model_components(file_path): | |
| # Ensure the file path is a Path object for robust path handling | |
| file_path = Path(file_path) | |
| # Check if the file exists | |
| if not file_path.is_file(): | |
| raise FileNotFoundError(f"The specified joblib file was not found: {file_path}") | |
| # Load all the model components from the joblib file | |
| model_components = joblib.load(file_path) | |
| # Create a reverse mapping from row index to tag | |
| if 'tag_to_row_index' in model_components: | |
| model_components['row_to_tag'] = {idx: tag for tag, idx in model_components['tag_to_row_index'].items()} | |
| return model_components | |
| # Load all components at the start | |
| tf_idf_components = load_model_components('tf_idf_files_420.joblib') | |
| idf = tf_idf_components['idf'] | |
| if isinstance(idf, dict): | |
| # idf is term -> idf_value; build a column-aligned vector | |
| t2c = tf_idf_components['tag_to_column_index'] | |
| n_cols = max(t2c.values()) + 1 | |
| idf_by_col = np.ones(n_cols, dtype=np.float32) | |
| for term, col in t2c.items(): | |
| idf_by_col[col] = float(idf.get(term, 1.0)) | |
| tf_idf_components['idf'] = idf_by_col | |
| 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) | |
| # Read the set of valid artists into memory. | |
| artist_set = set() | |
| with open("fluffyrock_3m.csv", 'r', newline='', encoding='utf-8') as csvfile: | |
| """ | |
| Load artist names from a CSV file and store them in the global set. | |
| Artist tags start with 'by_' and the prefix will be removed. | |
| """ | |
| reader = csv.reader(csvfile) | |
| for row in reader: | |
| tag_name = row[0] # Assuming the first column contains the tag names | |
| if tag_name.startswith('by_'): | |
| # Strip 'by_' from the start of the tag name and add to the set | |
| artist_name = tag_name[3:] # Remove the first three characters 'by_' | |
| artist_set.add(artist_name) | |
| def is_artist(name): | |
| return name in artist_set | |
| sample_images_directory_path = 'sampleimages' | |
| def generate_artist_image_tuples(top_artists, image_directory): | |
| json_files = glob.glob(f'{image_directory}/*.json') | |
| if not json_files: | |
| return [], [] # no mapping present; return empty galleries safely | |
| json_file_path = json_files[0] | |
| with open(json_file_path, 'r') as json_file: | |
| artist_to_file_map = json.load(json_file) | |
| # DEBUG: mapping + baseline info | |
| logging.debug("Gallery %s: loaded %d entries (map file=%s)", | |
| image_directory, len(artist_to_file_map), json_file_path) | |
| _base = artist_to_file_map.get("") | |
| logging.debug( | |
| "Gallery %s: baseline '' -> %r (exists=%s)", | |
| image_directory, | |
| _base, | |
| os.path.exists(os.path.join(image_directory, _base)) if _base else None, | |
| ) | |
| baseline_tuple = [] | |
| filename = artist_to_file_map.get("") | |
| if filename: | |
| image_path = os.path.join(image_directory, filename) | |
| if os.path.exists(image_path): | |
| baseline_tuple = [(image_path, "No Artist")] | |
| artist_image_tuples = [] | |
| for artist in top_artists: | |
| filename = artist_to_file_map.get(artist) | |
| # DEBUG: per-artist resolution | |
| logging.debug( | |
| "Gallery %s: %s -> %r (exists=%s)", | |
| image_directory, | |
| artist, | |
| filename, | |
| os.path.exists(os.path.join(image_directory, filename)) if filename else None, | |
| ) | |
| 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 baseline_tuple, 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 | |
| def create_html_tables_for_tags(subtable_heading, item_heading, word_similarity_tuples, 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;'><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>" | |
| # Loop through the results and add table rows for each | |
| for word, sim in word_similarity_tuples: | |
| word_with_underscores = word.replace(' ', '_') | |
| word_with_escaped_parentheses = escape_parens_for_display(word) | |
| count = tag2count.get(word_with_underscores.replace("\\(", "(").replace("\\)", ")"), 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 class='hover-underline' href='{wiki_url}' target='_blank' " | |
| f"title='{wiki_entry}'>{word_with_escaped_parentheses}</a>" | |
| ) | |
| else: | |
| # Display the word without any hyperlink or tooltip | |
| tag_element = word_with_escaped_parentheses | |
| # 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 create_top_artists_table(top_artists): | |
| html_str = "<div class=\"scrollable-content\" style='display: inline-block; margin: 20px; text-align: center;'>" | |
| html_str += "<h1>Top Artists</h1>" | |
| html_str += "<table style='font-size: 20px; border-collapse: collapse;'>" | |
| html_str += "<thead><tr><th>Artist</th><th>Similarity</th></tr></thead><tbody>" | |
| for artist, score in top_artists: | |
| artist_disp = escape_html(escape_parens_for_display(artist)) | |
| similarity_percentage = "{:.1f}%".format(score * 100) | |
| html_str += ( | |
| f"<tr><td style='padding: 3px 20px; border: none;'>{artist_disp}</td>" | |
| f"<td style='padding: 3px 20px; border: none;'>{similarity_percentage}</td></tr>" | |
| ) | |
| html_str += "</tbody></table></div>" | |
| return html_str | |
| def construct_pseudo_vector(pseudo_doc_terms, idf, term_to_column_index): | |
| cols, data = [], [] | |
| for term, w in pseudo_doc_terms.items(): | |
| j = term_to_column_index.get(term) | |
| if j is None: | |
| continue | |
| cols.append(j) | |
| data.append(w * idf[j]) | |
| n_cols = len(idf) | |
| indptr = [0, len(cols)] | |
| return csr_matrix((data, cols, indptr), shape=(1, n_cols), dtype=np.float32) | |
| def get_top_indices(reduced_pseudo_vector, reduced_matrix): | |
| # Compute cosine similarities | |
| similarities = cosine_similarity(reduced_pseudo_vector, reduced_matrix).flatten() | |
| # Get sorted tag indices based on similarities, in descending order | |
| sorted_indices = np.argsort(-similarities) | |
| # Return the top N indices | |
| return sorted_indices | |
| def get_tfidf_reduced_similar_tags(pseudo_doc_terms, allow_nsfw_tags): | |
| idf = tf_idf_components['idf'] | |
| term_to_column_index = tf_idf_components['tag_to_column_index'] | |
| row_to_tag = tf_idf_components['row_to_tag'] | |
| svd = tf_idf_components['svd_model'] | |
| # 1) Build the pseudo TF-IDF, reduce to SVD space (unchanged) | |
| pseudo_tfidf_vector = construct_pseudo_vector(pseudo_doc_terms, idf, term_to_column_index) | |
| reduced_pseudo_vector = svd.transform(pseudo_tfidf_vector) # shape (1, D) | |
| # 2) ANN: only fetch nearest non-artist candidates (no full-matrix cosine) | |
| K = 2000 # tune for speed/recall | |
| top_inds, top_sims = _ann_tags_topk(reduced_pseudo_vector, k=K) | |
| # 3) Build similarity dict from those candidates | |
| tag_similarity_dict = {} | |
| for i, sim in zip(top_inds, top_sims): | |
| tag = row_to_tag.get(int(i)) | |
| if tag is not None: | |
| tag_similarity_dict[tag] = float(sim) | |
| if not allow_nsfw_tags: | |
| tag_similarity_dict = {t: s for t, s in tag_similarity_dict.items() if t not in nsfw_tags} | |
| # 4) Sort & escape like before | |
| sorted_tag_similarity_dict = OrderedDict(sorted(tag_similarity_dict.items(), key=lambda x: x[1], reverse=True)) | |
| transformed_sorted_tag_similarity_dict = OrderedDict( | |
| (key.replace('_', ' ').replace('(', '\\(').replace(')', '\\)'), val) | |
| for key, val in sorted_tag_similarity_dict.items() | |
| ) | |
| return transformed_sorted_tag_similarity_dict | |
| def create_html_placeholder(title="", content="", placeholder_height=400, placeholder_width="100%"): | |
| # Include a title in the same style as the top artists table heading | |
| html_placeholder = f"<div class=\"scrollable-content\" style='text-align: center;'><h1>{title}</h1></div>" | |
| # Conditionally add content if present | |
| if content: | |
| html_placeholder += f"<div style='text-align: center; margin-bottom: 20px;'><p>{content}</p></div>" | |
| # Add the placeholder div with specified height and width | |
| html_placeholder += f"<div style='height: {placeholder_height}px; width: {placeholder_width}; margin: 20px auto; background: transparent;'></div>" | |
| return html_placeholder | |
| def find_similar_tags(test_tags, tag_to_context_similarity, context_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() | |
| modified_tags = [tag_info['modified_tag'] for tag_info in test_tags] | |
| transformed_tags = [tag.replace(' ', '_') for tag in modified_tags] | |
| # Find similar tags and prepare data for tables | |
| html_content = "<div class=\"scrollable-content\" style='display: inline-block; margin: 20px; text-align: center;'>" | |
| html_content += "<h1>Unknown Tags</h1>" # Heading for the table | |
| tags_added = False | |
| bad_entities = [] | |
| known_entities_in_prompt = [] | |
| encountered_modified_tags = set() | |
| for tag_info in test_tags: | |
| original_tag = tag_info['original_tag'] | |
| modified_tag = tag_info['modified_tag'] | |
| start_pos = tag_info['start_pos'] | |
| end_pos = tag_info['end_pos'] | |
| node_type = tag_info['node_type'] | |
| # Build the underscore form up-front | |
| modified_tag_for_search = modified_tag.replace(' ', '_') | |
| if modified_tag in special_tags: | |
| bad_entities.append({"entity":"Special", "start":start_pos, "end":end_pos}) | |
| continue | |
| # Only accept exact underscore model-specific tokens (e.g., "score_9") | |
| # special score/rating tags (kept as-is) | |
| if modified_tag in special_tags: | |
| bad_entities.append({"entity": "Special", "start": start_pos, "end": end_pos}) | |
| continue | |
| # Model-specific tokens must match the user's input *exactly* (no pre-normalization). | |
| # Use the original token as typed in the prompt, lowercased. | |
| original_raw = tag_info["original_tag"].strip().lower() | |
| if original_raw in MODEL_SPECIFIC_TAGS: | |
| bad_entities.append({"entity": "Model Specific", "start": start_pos, "end": end_pos}) | |
| continue | |
| if modified_tag in encountered_modified_tags: | |
| bad_entities.append({"entity":"Duplicate", "start":start_pos, "end":end_pos}) | |
| continue | |
| encountered_modified_tags.add(modified_tag) | |
| norm_artist = ( | |
| modified_tag_for_search | |
| .lower() | |
| .removeprefix('by_') # tolerate users typing "by something" or not | |
| ) | |
| if is_artist(norm_artist): | |
| by_key = f"by_{norm_artist}" | |
| # try by_* first, then raw form as fallback | |
| count = (find_similar_tags.tag2count.get(by_key) or | |
| find_similar_tags.tag2count.get(modified_tag_for_search, 0)) | |
| tag_id, wiki_entry = ( | |
| find_similar_tags.tag2idwiki.get(by_key) or | |
| find_similar_tags.tag2idwiki.get(modified_tag_for_search, (None, '')) | |
| ) | |
| wiki_url = f"https://e621.net/wiki_pages/{tag_id}" if tag_id is not None and wiki_entry else "" | |
| known_entities_in_prompt.append({ | |
| "entity": "Known Tag", | |
| "start": start_pos, | |
| "end": end_pos, | |
| "count": count, | |
| "wiki_url": wiki_url, | |
| "wiki_entry": wiki_entry | |
| }) | |
| continue | |
| 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 modified_tag in find_similar_tags.tag2aliases and "_" in modified_tag: #Implicitly tell the user that they should get rid of the underscore | |
| result.append((modified_tag_for_search.replace('_',' '), 1)) | |
| seen.add(modified_tag) | |
| else: #The user correctly did not put underscores in their tag | |
| count = find_similar_tags.tag2count.get(modified_tag_for_search, 0) # Get the count if available, otherwise default to 0 | |
| tag_id, wiki_entry = find_similar_tags.tag2idwiki.get(modified_tag_for_search, (None, '')) | |
| # Check if tag_id and wiki_entry are valid | |
| wiki_url = "" | |
| 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}" | |
| known_entities_in_prompt.append({"entity":"Known Tag", "start":start_pos, "end":end_pos, "count":count, "wiki_url":wiki_url, "wiki_entry":wiki_entry}) | |
| 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 = [(w, s) for (w, s) in result if w.replace(' ', '_') not in nsfw_tags] | |
| # --- Context re-scoring (keys match how get_tfidf_reduced_similar_tags formats them) --- | |
| def _ctx_score(name: str) -> float: | |
| v = tag_to_context_similarity.get(name) | |
| if v is None: | |
| # TF-IDF dict escapes parentheses; candidates from FT do not. | |
| v = tag_to_context_similarity.get(name.replace('(', '\\(').replace(')', '\\)')) | |
| return float(v) if v is not None else 0.0 | |
| # If the slider is at 1.0, only keep candidates that exist in the TF-IDF context list. | |
| if context_similarity_weight >= 0.999: | |
| ctx_keys = set(tag_to_context_similarity.keys()) | |
| result = [ | |
| (w, s) for (w, s) in result | |
| if (w in ctx_keys) or (w.replace('(', '\\(').replace(')', '\\)') in ctx_keys) | |
| ] | |
| # Linear blend: final = (1-λ)*fasttext + λ*context (no extra 0.5 scaling) | |
| result = [ | |
| (w, (1.0 - context_similarity_weight) * s + context_similarity_weight * _ctx_score(w)) | |
| for (w, s) in result | |
| ] | |
| result = sorted(result, key=lambda x: x[1], reverse=True)[:20] | |
| html_content += create_html_tables_for_tags(modified_tag, "Corrected Tag", result, find_similar_tags.tag2count, find_similar_tags.tag2idwiki) | |
| bad_entities.append({"entity":"Unknown Tag", "start":start_pos, "end":end_pos}) | |
| tags_added=True | |
| # If no tags were processed, add a message; otherwise close the wrapper div | |
| if not tags_added: | |
| html_content = create_html_placeholder(title="Unknown Tags", content="No Unknown Tags Found") | |
| else: | |
| html_content += "</div>" | |
| return html_content, bad_entities, known_entities_in_prompt | |
| def build_tag_offsets_dicts(new_image_tags_with_positions): | |
| # Structure the data for HighlightedText | |
| tag_data = [] | |
| for tag_text, start_pos, nodetype in new_image_tags_with_positions: | |
| # Modify the tag | |
| modified_tag = tag_text.replace('_', ' ').replace('\\(', '(').replace('\\)', ')').strip() | |
| artist_matrix_tag = tag_text.replace('_', ' ').replace('\\(', '\(').replace('\\)', '\)').strip() | |
| tf_idf_matrix_tag = re.sub(r'\\([()])', r'\1', re.sub(r' ', '_', tag_text.strip().removeprefix('by ').removeprefix('by_'))) | |
| # Calculate the end position based on the original tag length | |
| end_pos = start_pos + len(tag_text) | |
| # Append the structured data for each tag | |
| tag_data.append({ | |
| "original_tag": tag_text, | |
| "start_pos": start_pos, | |
| "end_pos": end_pos, | |
| "modified_tag": modified_tag, | |
| "artist_matrix_tag": artist_matrix_tag, | |
| "tf_idf_matrix_tag": tf_idf_matrix_tag, | |
| "node_type": nodetype | |
| }) | |
| return tag_data | |
| def augment_bad_entities_with_regex(text): | |
| bad_entities = [] | |
| #comma at end | |
| match = re.search(r',(?=\s*$)', text) | |
| if match: | |
| index = match.start() | |
| bad_entities.append({"entity":"Remove Final Comma", "start":index, "end":index+1}) | |
| match = re.search(r'\([^()]*(,)\s*\)\s*$', text) | |
| if match: | |
| index = match.start(1) | |
| bad_entities.append({"entity":"Remove Final Comma", "start":index, "end":index+1}) | |
| match = re.search(r'\([^()]*(,)\s*:\s*\d+(\.\d+)?\s*\)\s*$', text) | |
| if match: | |
| index = match.start(1) | |
| bad_entities.append({"entity":"Remove Final Comma", "start":index, "end":index+1}) | |
| # Comma after parentheses, multiple occurrences | |
| for match in re.finditer(r'(?<!\\)\)\s*(,)\s*[^\s]', text): | |
| index = match.start(1) | |
| bad_entities.append({"entity": "Move Comma Inside Parentheses", "start": index, "end": index + 1}) | |
| # Double Comma detection | |
| for match in re.finditer(r',\s*,', text): | |
| index = match.start() | |
| bad_entities.append({"entity": "Double Comma", "start": index, "end": index + match.end() - match.start()}) | |
| return bad_entities | |
| def escape_html(text): | |
| return text.replace("&", "&").replace("<", "<").replace(">", ">").replace('"', """).replace("'", "'") | |
| def escape_parens_for_display(s: str) -> str: | |
| # ensure single backslash before any literal parens in display text | |
| return ( | |
| s.replace("\\(", "(") | |
| .replace("\\)", ")") | |
| .replace("(", "\\(") | |
| .replace(")", "\\)") | |
| ) | |
| def format_annotated_html(bad_entities, known_entities, text): | |
| tooltip_map = { | |
| "Unknown Tag": "This may not be a valid e621 tag. Consider removing or replacing it with tag(s) from the \"Unknown Tags\" section.", | |
| "Duplicate": "This tag has appeared multiple times in your prompt. Consider removing the copies.", | |
| "Remove Final Comma": "There should be no comma at the end of your prompt. Consider removing it.", | |
| "Move Comma Inside Parentheses": "In most e621-based models, the comma following a tag functions as an "attention anchor", carrying most of the tag's information. It should therefore be assigned the same weight as the rest of the tag. So instead of "(lineless:1.1),", consider "(lineless,:1.1)" or "(lineless,)"", | |
| "Double Comma": "One comma between tags is considered ample.", | |
| "Model Specific": "This is not an e621 tag, but may still be valid with the right model. Check your model's documentation. If the tag is not mentioned in the documentation, do not use it." | |
| } | |
| color_map = { | |
| "Double Comma": ("black", "#E69F00"), # Orange — black text | |
| "Duplicate": ("black", "#F0E442"), # Yellow — black text | |
| "Model Specific": ("black", "#56B4E9"), # Sky blue — black text | |
| "Move Comma Inside Parentheses": ("white", "#009E73"), # Bluish green | |
| "Remove Final Comma": ("white", "#0072B2"), # Blue | |
| "Unknown Tag": ("white", "#D55E00"), # Vermilion — danger | |
| } | |
| # Splice from the original text so indexes stay valid. | |
| combined = sorted(bad_entities + known_entities, key=lambda x: x["start"], reverse=True) | |
| html_text = text | |
| for entity in combined: | |
| start = entity["start"] | |
| end = entity["end"] | |
| label = entity["entity"] | |
| # Escape only the replaced segment (keeps indices correct). | |
| segment = text[start:end] | |
| disp = escape_html(escape_parens_for_display(segment)) | |
| if label == "Known Tag": | |
| wiki_url = entity.get("wiki_url", "") | |
| count = entity.get("count", 0) | |
| wiki_entry = entity.get("wiki_entry", "") | |
| sanitized_wiki_entry = escape_html(wiki_entry) if wiki_entry else "Unavailable" | |
| if wiki_url: | |
| html_part = ( | |
| f'<a class="hover-underline" href="{wiki_url}" target="_blank" ' | |
| f'title="Count: {count}\tWiki: {sanitized_wiki_entry}" ' | |
| f'style="cursor: pointer; font-style: italic;">{disp}</a>' | |
| ) | |
| else: | |
| html_part = ( | |
| f'<span class="hover-underline" title="Count: {count}\tWiki: {sanitized_wiki_entry}" ' | |
| f'style="cursor: help; font-style: italic;">{disp}</span>' | |
| ) | |
| else: | |
| fg, bg = color_map.get(label, ("black", "white")) | |
| # Make single-character issues (commas) easier to see by padding the span | |
| pad_labels = ("Remove Final Comma", "Move Comma Inside Parentheses", "Double Comma") | |
| pad_class = "badpad" if label in pad_labels else "" | |
| class_attr = f' class="{pad_class}"' if pad_class else "" | |
| html_part = f'<span{class_attr} style="background-color: {bg}; color: {fg};">{disp}</span>' | |
| html_text = html_text[:start] + html_part + html_text[end:] | |
| # Color key (only for labels that actually appeared) | |
| color_key_html = "<div style='text-align: right; margin-top: 20px;'>Key:" | |
| used_labels = {e["entity"] for e in bad_entities} | |
| for label, (fg, bg) in color_map.items(): | |
| if label in used_labels: | |
| tooltip = tooltip_map.get(label, "") | |
| color_key_html += ( | |
| f" <span class='hover-underline' style='background-color: {bg}; color: {fg}; margin-right: 10px;' " | |
| f"title='{tooltip}'>{label}</span>" | |
| ) | |
| color_key_html += "</div>" | |
| # Wrap the whole annotated area so we can place the hint inside it | |
| annotated_box = ( | |
| "<div class='annotated-wrap' style='padding:10px;font-size:16px;'>" | |
| f"{html_text}" | |
| f"{TOOLTIP_NOTE_HTML}" | |
| "</div>" | |
| ) | |
| return annotated_box + color_key_html | |
| def find_similar_artists(original_tags_string, top_n, context_similarity_weight, allow_nsfw_tags): | |
| try: | |
| new_tags_string = original_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) | |
| tag_data = build_tag_offsets_dicts(new_image_tags) | |
| #Suggested tags stuff | |
| suggested_tags_html_content = "<div class=\"scrollable-content\" style='display: inline-block; margin: 20px; text-align: center;'>" | |
| suggested_tags_html_content += "<h1>Suggested Tags</h1>" # Heading for the table | |
| terms = [item["tf_idf_matrix_tag"] for item in tag_data] + removed_tags | |
| suggested_tags = get_tfidf_reduced_similar_tags(dict(Counter(terms)), allow_nsfw_tags) | |
| unseen_tags_data, bad_entities, known_entities = find_similar_tags(tag_data, suggested_tags, context_similarity_weight, allow_nsfw_tags) | |
| #Bad tags stuff | |
| bad_entities.extend(augment_bad_entities_with_regex(new_tags_string)) | |
| bad_entities.sort(key=lambda x: x['start']) | |
| #bad_tags_illustrated_string = {"text":new_tags_string, "entities":bad_entities} | |
| bad_tags_illustrated_html = format_annotated_html(bad_entities, known_entities, new_tags_string) | |
| # Create a set of tags that should be filtered out | |
| filter_tags = {entry["original_tag"].strip() for entry in tag_data} | |
| filter_tags_norm = { _norm_tag_for_lookup(t.lower().removeprefix('by ').removeprefix('by_')) for t in filter_tags } | |
| suggested_tags_filtered = OrderedDict( | |
| (k, v) for k, v in suggested_tags.items() | |
| if k not in filter_tags and _norm_tag_for_lookup(k.lower()) not in filter_tags_norm | |
| ) | |
| # Splitting the dictionary into two based on the condition | |
| def _norm_no_by(s: str) -> str: | |
| n = _norm_tag_for_lookup(s) | |
| return n[3:] if n.startswith("by_") else n | |
| suggested_artist_tags_filtered = OrderedDict( | |
| (k, v) for k, v in suggested_tags_filtered.items() | |
| if is_artist(_norm_no_by(k)) | |
| ) | |
| suggested_non_artist_tags_filtered = OrderedDict( | |
| (k, v) for k, v in suggested_tags_filtered.items() | |
| if not is_artist(_norm_no_by(k)) and k not in special_tags | |
| ) | |
| topnsuggestions = list(islice(suggested_non_artist_tags_filtered.items(), 100)) | |
| suggested_tags_html_content += create_html_tables_for_tags("-", "Suggested Tag", topnsuggestions, find_similar_tags.tag2count, find_similar_tags.tag2idwiki) | |
| suggested_tags_html_content += "</div>" | |
| # --- Artist stuff: query artist-only index directly --- | |
| idf_vec = tf_idf_components['idf'] | |
| t2c = tf_idf_components['tag_to_column_index'] | |
| svd = tf_idf_components['svd_model'] | |
| pseudo_terms = dict(Counter(terms)) | |
| pseudo_vec = construct_pseudo_vector(pseudo_terms, idf_vec, t2c) | |
| reduced_q = svd.transform(pseudo_vec) | |
| K_art = max(100, top_n * 10) # widen search to stabilize ranks | |
| art_inds, art_sims = _ann_artists_topk(reduced_q, k=K_art) | |
| row_to_tag = tf_idf_components['row_to_tag'] | |
| bad_labels = {"by_unknown_artist", "by_conditional_dnp", "unknown_artist", "conditional_dnp"} | |
| top_artists_raw = [] | |
| for idx_i, sim in zip(art_inds, art_sims): | |
| tag = row_to_tag.get(int(idx_i), "") | |
| if not tag: | |
| continue | |
| # Normalize spaces to underscores for reliable checks | |
| norm = tag.replace(" ", "_") | |
| # Drop known non-artist placeholders | |
| if norm in bad_labels: | |
| continue | |
| # Accept either "by_foo" or plain "foo" | |
| base = norm[3:] if norm.startswith("by_") else norm | |
| # Guard: only keep if this *really* is an artist we know | |
| if not is_artist(base): | |
| continue | |
| name_disp = base.replace("_", " ") | |
| top_artists_raw.append((name_disp, float(sim))) | |
| if not top_artists_raw: | |
| logging.debug("No artist hits. First few neighbor labels: %s", | |
| [row_to_tag.get(int(i), "") for i in art_inds[:10]]) | |
| # take the best unique names, in order | |
| seen = set() | |
| deduped = [] | |
| for n, s in top_artists_raw: | |
| if n not in seen: | |
| deduped.append((n, s)) | |
| seen.add(n) | |
| if len(deduped) >= top_n: | |
| break | |
| top_artists = deduped | |
| logging.debug("Top artists (n=%d): %s", len(top_artists), top_artists) | |
| top_artists_str = create_top_artists_table(top_artists) | |
| dynamic_prompts_formatted_artists = "{" + "|".join([artist for artist, _ in top_artists]) + "}" | |
| dynamic_prompts_formatted_artists = "{" + "|".join( | |
| [escape_parens_for_display(artist) for artist, _ in top_artists] | |
| ) + "}" | |
| image_galleries = [] | |
| for root, dirs, files in os.walk(sample_images_directory_path): | |
| for name in dirs: | |
| baseline, artists = generate_artist_image_tuples([name for name, _ in top_artists], os.path.join(root, name)) | |
| dir_path = os.path.join(root, name) | |
| baseline, artists = generate_artist_image_tuples([n for n, _ in top_artists], dir_path) | |
| logging.debug("Gallery built for %s -> baseline=%d, artists_found=%d", dir_path, len(baseline), len(artists)) | |
| image_galleries.append(baseline) # Add baseline as its own gallery item | |
| image_galleries.append(artists) # Extend the list with artist tuples | |
| return (unseen_tags_data, bad_tags_illustrated_html, suggested_tags_html_content, top_artists_str, dynamic_prompts_formatted_artists, *image_galleries) | |
| except ParseError: | |
| # Build empty galleries so the tuple length matches the declared outputs | |
| empty_galleries = [] | |
| for _root, _dirs, _files in os.walk(sample_images_directory_path): | |
| for _ in _dirs: | |
| empty_galleries.extend([[], []]) # one empty list per Gallery component | |
| return ( | |
| create_html_placeholder(title="Unknown Tags", content="Parse Error"), | |
| "Parse Error: Check for mismatched parentheses or something", | |
| create_html_placeholder(title="Suggested Tags"), | |
| "", # top_artists | |
| "", # dynamic_prompts | |
| *empty_galleries, | |
| ) | |
| with gr.Blocks(css=css) as app: | |
| with gr.Group(): | |
| with gr.Row(): | |
| with gr.Column(scale=3, elem_classes=["prompt-col"]): | |
| image_tags = gr.Textbox( | |
| label="Enter Prompt", | |
| placeholder="e.g. fox, outside, detailed background, ...", | |
| lines=1 # Enter submits (see .submit() below) | |
| ) | |
| bad_tags_illustrated_string = gr.HTML() | |
| with gr.Column(scale=1): | |
| # Load once so Gradio serves it directly (no /file= URL needed) | |
| _mascot_pil = Image.open(MASCOT_FILE).convert("RGBA") | |
| mascot_img = gr.Image( | |
| value=_mascot_pil, # pass the in-memory image | |
| show_label=False, | |
| interactive=False, | |
| height=220, # keep the same visual size | |
| elem_id="mascot" | |
| ) | |
| submit_button = gr.Button(variant="primary") | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| with gr.Group(): | |
| with gr.Row(): | |
| context_similarity_weight = gr.Slider(minimum=0, maximum=1, value=0.5, step=0.1, label="Context Similarity Weight") | |
| allow_nsfw = gr.Checkbox(label="Include Mature Tags", value=False) | |
| with gr.Row(): | |
| with gr.Column(scale=2, elem_classes=["pane-left"]): | |
| unseen_tags = gr.HTML( | |
| label="Unknown Tags", | |
| value=create_html_placeholder(title="Unknown Tags"), | |
| elem_id="unseen_html", | |
| elem_classes=["scroll-fade"], | |
| ) | |
| with gr.Column(scale=1, elem_classes=["pane-right"]): | |
| suggested_tags = gr.HTML( | |
| label="Suggested Tags", | |
| value=create_html_placeholder(title="Suggested Tags"), | |
| elem_id="suggested_html", | |
| elem_classes=["scroll-fade"], | |
| ) | |
| with gr.Column(scale=1): | |
| with gr.Group(): | |
| num_artists = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Number of artists") | |
| top_artists = gr.HTML( | |
| label="Top Artists", | |
| value=create_html_placeholder(title="Top Artists"), | |
| elem_id="artists_html", | |
| elem_classes=["scroll-fade"], | |
| ) | |
| gr.HTML(""" | |
| <script> | |
| (function(){ | |
| function wire(id){ | |
| const host = document.getElementById(id); | |
| if (!host) return; | |
| // Always use the *inner* .scrollable-content as the scroller | |
| const getScroller = () => host.querySelector('.scrollable-content') || host; | |
| let scroller = getScroller(); | |
| // Set CSS var so the fade blends with host background | |
| const bg = getComputedStyle(host).backgroundColor; | |
| host.style.setProperty('--host-bg', bg); | |
| const refresh = () => { | |
| // guard for fractional pixels across browsers | |
| const atBottom = Math.ceil(scroller.scrollTop + scroller.clientHeight) >= scroller.scrollHeight; | |
| host.classList.toggle('at-bottom', atBottom); | |
| }; | |
| // (Re)attach scroll listener to the current scroller | |
| const attach = (el) => { | |
| if (!el) return; | |
| el.addEventListener('scroll', refresh, {passive:true}); | |
| // initial state | |
| refresh(); | |
| }; | |
| attach(scroller); | |
| // If Gradio replaces inner HTML, re-wire to new scroller | |
| new MutationObserver(() => { | |
| const next = getScroller(); | |
| if (next && next !== scroller) { | |
| scroller.removeEventListener && scroller.removeEventListener('scroll', refresh); | |
| scroller = next; | |
| attach(scroller); | |
| } | |
| // background might change with themes; keep it fresh | |
| const newBg = getComputedStyle(host).backgroundColor; | |
| host.style.setProperty('--host-bg', newBg); | |
| refresh(); | |
| }).observe(host, {childList: true, subtree: true}); | |
| // Also respond to resizes | |
| new ResizeObserver(refresh).observe(host); | |
| } | |
| ['unseen_html','suggested_html','artists_html'].forEach(wire); | |
| })(); | |
| </script> | |
| """, visible=False) | |
| 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.") | |
| galleries = [] | |
| for root, dirs, files in os.walk(sample_images_directory_path): | |
| for name in dirs: | |
| with gr.Row(): | |
| baseline = gr.Gallery(allow_preview=False, rows=1, columns=1, height=420, scale=3) | |
| styles = gr.Gallery(allow_preview=False, rows=2, columns=5, height=420, scale=8) | |
| galleries.extend([baseline, styles]) | |
| submit_button.click( | |
| find_similar_artists, | |
| inputs=[image_tags, num_artists, context_similarity_weight, allow_nsfw], | |
| outputs=[unseen_tags, bad_tags_illustrated_string, suggested_tags, top_artists, dynamic_prompts] + galleries | |
| ) | |
| # Also run when pressing Enter in the prompt box | |
| image_tags.submit( | |
| find_similar_artists, | |
| inputs=[image_tags, num_artists, context_similarity_weight, allow_nsfw], | |
| outputs=[unseen_tags, bad_tags_illustrated_string, suggested_tags, top_artists, dynamic_prompts] + galleries | |
| ) | |
| gr.Markdown(faq_content) | |
| app.queue().launch(allowed_paths=[str(MASCOT_DIR)]) | |