Add artists and "de-aligned" tags to NAI's tag suggestions

README.md

NAI Tags+

This userscript injects artists and tags that do not conform to "AI alignment" standards into NovelAI's suggested imagegen tags.

Prerequisites

You will need a userscript extension. Below is a suggestion.

• https://violentmonkey.github.io

Installation

Open this URL: https://gist.github.com/catboxanon/cfcd89f2838b410d3c5b10ca46e799da/raw/nai-tags-plus.user.js

If you have an open tab on NovelAI, refresh it.

If you do not recieve an install dialog, you do not have a userscript extension, or it is not functioning properly.

Usage

In the prompt area, enter tags to retrieve tag suggestions as normal (assuming tag suggestions are enabled). Tags suggested by this userscript will be added to the bottom of the list.

To suggest artist tags, the tag must be prefixed by artist:

Changelog

(2024-07-20) 0.1.3 - Ignore invalid URL requests from causing errors

(2024-07-11) 0.1.2 - Alias SFW tags that NAI renamed.

(2024-07-11) 0.1.1 - Add support for aliasing certain tags (due to NAI renaming some tags).

(2024-07-11) 0.1.0 - Initial version

nai-tags-plus.user.js

// ==UserScript==
// @name        NAI Tags+
// @namespace   hdg-nai-tags-plus
// @match       https://novelai.net/*
// @grant       none
// @version     0.1.3
// @author      Anonymous
// @description Suggest artists and "de-aligned" tags in NAI
// @updateURL   https://gist.github.com/catboxanon/cfcd89f2838b410d3c5b10ca46e799da/raw/nai-tags-plus.user.js
// @downloadURL https://gist.github.com/catboxanon/cfcd89f2838b410d3c5b10ca46e799da/raw/nai-tags-plus.user.js
// ==/UserScript==

const DATA_URL = 'https://raw.githubusercontent.com/DominikDoom/a1111-sd-webui-tagcomplete/1d8d9f64b5aaaa0b2b9cbebaa5fd45a48b193ba6/tags/danbooru.csv';
const TAG_MAP = {
  "pubic_tattoo": "womb_tattoo",
  "v": "peace_sign",
  "double_v": "double_peace",
  "|_|": "bar_eyes",
  "\\||\/": "open_\/m\/",
  ":|": "neutral_face",
  ";|": "neutral_face",
};

let generalData;
let artistsData;

async function fetchCSV(url) {
  const response = await window.fetch(url);
  const text = await response.text();
  return parseCSV(text);
}

function parseCSV(text) {
  const rows = text.split('\n').filter(row => row.trim() !== '');
  return rows.map(parseCSVRow).filter(row => (row[1] === '1' || row[1] === '0') && row[0] !== 'banned_artist');
}

function parseCSVRow(row) {
  const result = [];
  let inQuotes = false;
  let value = '';

  for (let i = 0; i < row.length; i++) {
    const char = row[i];

    if (char === '"' && (i === 0 || row[i - 1] !== '\\')) {
      inQuotes = !inQuotes;
    } else if (char === ',' && !inQuotes) {
      result.push(value.trim());
      value = '';
    } else {
      value += char;
    }
  }

  result.push(value.trim());

  return result;
}

function interpolate(value, min, max, newMin, newMax) {
  return (value - min) / (max - min) * (newMax - newMin) + newMin;
}

function preprocessData(data, type) {
  const preprocessedData = data
    .filter(row => parseInt(row[1]) === type)
    .map(row => {
    let tag = row[0];
    let aliases = row[3];
    if (tag in TAG_MAP) {
      aliases += ',' + tag;
      tag = TAG_MAP[tag];
    }
    return {
      tag: tag.replace(/_/g, ' '),
      count: parseInt(row[2], 10),
      aliases: aliases.replace(/_/g, ' ').split(',').filter(Boolean).map(alias => alias.trim())
    };
  });

  const capCount = 1500;
  const maxCount = Math.min((type === 1 ? capCount : 1_000_000_000), Math.max(...preprocessedData.map(artist => artist.count)));
  const minCount = Math.min(...preprocessedData.map(artist => artist.count));

  return preprocessedData.map(tag => ({
    ...tag,
    confidence: Math.min(tag.count, capCount) / maxCount,
    interpolatedCount: type === 1 ? interpolate(tag.count, minCount, maxCount, 0, 10000) : tag.count
  })).sort((a, b) => b.count - a.count || a.tag.localeCompare(b.tag));
}

function searchTags(query, data, prefix = '') {
  return data.filter(tag =>
    (
    tag.tag.includes(query)
    || tag.aliases.some(alias => alias.includes(query))
    || tag.tag.replace(/\s/g, '_').includes(query)
    || tag.aliases.some(alias => alias.replace(/\s/g, '_').includes(query))
    ) && !(/^[^\s]{1,2}\s[^\s]{1,2}$/).test(tag.tag)  // Exclude returning tags like "@ @" that NAI formats as "@_@"
  ).slice(0, 10).map(tag => ({
    tag: prefix + tag.tag,
    count: tag.interpolatedCount,
    confidence: tag.confidence
  }));
}

(async function initialize() {
  const rawData = await fetchCSV(DATA_URL);
  generalData = preprocessData(rawData, type = 0);
  artistsData = preprocessData(rawData, type = 1);
  // console.debug('Filtered and processed tag data: ', artistsData, generalData);
  main();
})();

function main() {
  window.nativeFetch = window.fetch;

  window.customFetch = async function(request, headers) {
    let req;
    let response;

    if (typeof request === 'string') {
      let url;
      let validUrl;
      try {
        url = new URL(request);
        validUrl = true;
      } catch {
        validUrl = false;
      }

      if (
        validUrl
        && url.hostname === 'image.novelai.net'
        && url.pathname === '/ai/generate-image/suggest-tags'
      ) {
        const promptSearchParam = url.searchParams.get('prompt');
        const isArtistQuery = promptSearchParam?.startsWith('artist:');
        const promptQuery = isArtistQuery ? promptSearchParam.substring(promptSearchParam.indexOf(':') + 1).trim() : promptSearchParam;

        const searchResults = searchTags(promptQuery, isArtistQuery ? artistsData : generalData, prefix = (isArtistQuery ? 'artist:' : ''));
        // console.debug('Search results:', searchResults);

        req = new Request(request, headers);
        response = await window.nativeFetch(req);
        const clonedResponse = response.clone();
        const responseBody = await clonedResponse.json();

        if (isArtistQuery) {
          responseBody.tags = searchResults;
        } else {
          const existingTags = new Set(responseBody.tags.map(tag => tag.tag));
          const filteredResults = searchResults.filter(result => !existingTags.has(result.tag));
          responseBody.tags = [...responseBody.tags, ...filteredResults];
        }

        const modifiedResponse = new Response(JSON.stringify(responseBody), {
          status: response.status,
          statusText: response.statusText,
          headers: response.headers
        });

        return modifiedResponse;
      }

      req = new Request(request, headers);
      response = await window.nativeFetch(req);
      response.requestInputObject = req;
    } else {
      response = await window.nativeFetch(request, headers);
    }

    if (typeof request === 'object') {
      response.requestInputObject = request;
    } else {
      response.requestInputURL = request;
      response.requestInputObject = req;
    }

    if (headers) {
      response.requestInputHeaders = headers;
    }

    return response;
  }

  window.fetch = window.customFetch;
}