|
// ==UserScript== |
|
// @name NAI Tags+ |
|
// @namespace hdg-nai-tags-plus |
|
// @match https://novelai.net/* |
|
// @grant none |
|
// @version 0.2.0 |
|
// @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/40ad070a02033bdd159214eee87531701cc43ef2/tags/danbooru_e621_merged.csv'; |
|
const TAG_MAP = { |
|
"pubic_tattoo": "womb_tattoo", |
|
"v": "peace_sign", |
|
"double_v": "double_peace", |
|
"|_|": "bar_eyes", |
|
"\\||\/": "open_\/m\/", |
|
":|": "neutral_face", |
|
";|": "neutral_face", |
|
"eyepatch_bikini": "square_bikini", |
|
"tachi-e": "character image", |
|
}; |
|
|
|
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; |
|
} |