diff --git a/README.md b/README.md index 4c164fe..5cf2ca4 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,11 @@ # Prompt Metadata Checker for Stable Diffusion Images -Reads the metadata of an AI generated image (using the `Exif` or `tEXT` enclosed in that image file). Populates the prompt, seed, sampler, and more. Known to work well with images generated using AUTOMATIC1111, InvokeAI or NMKD. Supports both `.png` and `.jpeg` images. - -If you run into any issues: please open up an Issue here on Github or talk to us on Discord. +Reads the metadata of an AI generated image (using the `Exif` or `tEXT` enclosed in that image file). Populates the prompt, seed, sampler, and more. Known to work well with images generated using AUTOMATIC1111, InvokeAI or NMKD. Supports both `.jpeg` and `.png` images. This is a port into PromptHero of a community extension originally made by [@HE1CO](https://github.com/HE1CO/Postie) and later revamped by [@drhino](https://github.com/drhino) +If you run into any issues: please open up an Issue here on Github or [talk to us on Discord](https://discord.com/channels/1026222136790110259) 👋 + ## Motivation This is what we use in [our prompt upload page](https://prompthero.com/prompt/upload). When you drop a valid image generated with A1111 or a supported program, this is what auto-fills all the fields in the page. diff --git a/src/controllers/prompts/upload_controller.js b/src/controllers/prompts/upload_controller.js index 36d25fb..ecb7719 100644 --- a/src/controllers/prompts/upload_controller.js +++ b/src/controllers/prompts/upload_controller.js @@ -106,21 +106,21 @@ export default class extends Controller { // Waits until the AJAX is done onDone.then( e => { modelVers.value = promptInfo.model_used_version - }); - /* - if (modelSlug.value) - modelSlug.setAttribute('disabled', true) - else - modelSlug.value = '' - - if (modelVers.value) - modelVers.setAttribute('disabled', true) - else - modelVers.value = '' - */ - if (!modelSlug.value) modelSlug.value = '' - if (!modelVers.value) modelVers.value = '' + /* + if (modelSlug.value) + modelSlug.setAttribute('disabled', true) + else + modelSlug.value = '' + + if (modelVers.value) + modelVers.setAttribute('disabled', true) + else + modelVers.value = '' + */ + if (!modelSlug.value) modelSlug.value = '' + if (!modelVers.value) modelVers.value = '' + }); } } \ No newline at end of file diff --git a/src/vendor/postie/dataEnricher.js b/src/vendor/postie/dataEnricher.js index 2bedeaa..a145e62 100644 --- a/src/vendor/postie/dataEnricher.js +++ b/src/vendor/postie/dataEnricher.js @@ -1,10 +1,23 @@ import hashes from "vendor/postie/modelHashes" +/** + * Returns the model and version from the model_hash + * The returned values are fixed identifiers from PromptHero.com + * @param {string} hash + * @returns { { model: string|undefined, version: string|undefined } } + */ const findModelByHash = hash => { if (hash) { + const hash_lowercase = hash.toLowerCase() + for (const [model, versions] of Object.entries(hashes)) { for (const [version, hashes] of Object.entries(versions)) { + if (hashes.find(x => x.toLowerCase() === hash_lowercase)) { + return { model, version } + } + + /* // Matches hash from Civit AI if (hashes.includes(hash.toUpperCase())) { return { model, version } @@ -14,6 +27,7 @@ const findModelByHash = hash => { if (hashes.includes(hash.toLowerCase())) { return { model, version } } + */ } } } @@ -23,9 +37,10 @@ const findModelByHash = hash => { /** + * Returns the normalized sampler name, if available on PromptHero.com * @param {string} sampler - * @returns {string} -*/ + * @returns {string|undefined} + */ const findSamplerByName = sampler => { // @TODO: are these mappings correct? and complete? const samplersDict = { @@ -52,25 +67,41 @@ const findSamplerByName = sampler => { "lms_karras": ["LMS Karras", "k_lms_karras"], } - return Object.keys(samplersDict).find( - x => x === sampler.toLowerCase() || samplersDict[x].includes(sampler) - ) ?? sampler + return Object.keys(samplersDict).find(x => samplersDict[x].includes(sampler)) } /** - * @param {string} hires_upscaler - * @returns {string} -*/ + * Returns the normalized upscaler, if available on PromptHero.com + * @param {string|undefined} hires_upscaler + * @returns {string|undefined} + */ const findUpscaler = hires_upscaler => { - if (hires_upscaler && hires_upscaler.indexOf('ESRGAN 4x') !== -1) { + if ( ! hires_upscaler ) { + return + } + + if (hires_upscaler.indexOf('ESRGAN 4x') !== -1) { return 'ESRGAN 4x' } - else if (hires_upscaler && hires_upscaler.indexOf('CodeFormer') !== -1) { + + if (hires_upscaler.indexOf('CodeFormer') !== -1) { return 'CodeFormer' - } else { - return '' } } -export { findModelByHash, findSamplerByName, findUpscaler }; \ No newline at end of file + +/** + * Generates the SHA-256 hash of a file + * @param {ArrayBuffer} arrayBuffer + * @returns {string} SHA-256 + * @see https://developer.mozilla.org/en-US/docs/Web/API/SubtleCrypto/digest + */ +const generateSha256Hash = async(arrayBuffer) => Array.from( + new Uint8Array( await crypto.subtle.digest('SHA-256', arrayBuffer) ) +).map( + x => x.toString(16).padStart(2, '0') +).join('') + + +export { findModelByHash, findSamplerByName, findUpscaler, generateSha256Hash }; diff --git a/src/vendor/postie/exifMetadataExtractor.js b/src/vendor/postie/exifMetadataExtractor.js index 7b84cc4..404c90a 100644 --- a/src/vendor/postie/exifMetadataExtractor.js +++ b/src/vendor/postie/exifMetadataExtractor.js @@ -22,7 +22,7 @@ const getTextFromImage = arrayBuffer => { return (uint8Array[start] << 8) | uint8Array[start + 1] case '32-bit': - return (uint8Array[start] << 24) | + return (uint8Array[start] << 24) | (uint8Array[start + 1] << 16) | (uint8Array[start + 2] << 8) | uint8Array[start + 3] diff --git a/src/vendor/postie/parseMetadata.js b/src/vendor/postie/parseMetadata.js index fc1e598..7489f23 100644 --- a/src/vendor/postie/parseMetadata.js +++ b/src/vendor/postie/parseMetadata.js @@ -1,12 +1,12 @@ import getTextFromImage from "vendor/postie/exifMetadataExtractor" -import { findSamplerByName, findUpscaler, findModelByHash } from "vendor/postie/dataEnricher" +import { findSamplerByName, findUpscaler, findModelByHash, generateSha256Hash } from "vendor/postie/dataEnricher" // Parses the PNG/JPEG metadata const parseMetadata = (file, callback) => { const fr = new FileReader() - fr.onload = () => { + fr.onload = async() => { let embed = getTextFromImage(fr.result) try { @@ -31,7 +31,9 @@ const parseMetadata = (file, callback) => { // Same as above but just the other way around embed.prompt = embed.prompt.match(/[^\[\]]+(?=\[(.*)\]|$)/g).join(', ') - embed.program_used = "InvokeAI" + if (embed.app_id === "invoke-ai/InvokeAI") { + embed.program_used = embed.model_hash === 0 ? "NMKD" : "InvokeAI" + } } catch (e) { // Rethrow unless parsing as JSON failed @@ -70,41 +72,44 @@ const parseMetadata = (file, callback) => { height, } + // @TODO: verify the strict use of this software embed.program_used = "AUTOMATIC1111" } // Normalization const normalizer = str => str?.trim() - .replace(/,( ?,)+/g, ',') // separates with a maximum of one comma - .replace(/^[\s?,]+|[\s?,]+$/g, '') // removes trailing/leading commas and whitespace - .replace(/\s+/g, ' ') // uses single spaces only + .replace(/(\s*,\s*)+/g, ', ') // separates with a maximum of one comma + .replace(/^[\s,]+|[\s,]+$/g, '') // removes trailing/leading commas and whitespace + .replace(/\s+/g, ' ') // uses single spaces only embed.prompt = normalizer(embed.prompt) embed.negative_prompt = normalizer(embed.negative_prompt) // Handles the matching website - console.log("Metadata parsed!", embed); - embed = enrichMetadataForPromptHero(embed); + console.log("parseMetadata.js: Metadata parsed!", embed); + embed = await enrichMetadataForPromptHero(embed, fr.result); callback(embed); } fr.readAsArrayBuffer(file) } -const enrichMetadataForPromptHero = promptInfo => { - var richPromptInfo = promptInfo; +const enrichMetadataForPromptHero = async(promptInfo, arrayBuffer) => { + const richPromptInfo = promptInfo; richPromptInfo.sampler_raw = promptInfo.sampler; richPromptInfo.sampler = findSamplerByName(promptInfo.sampler); - richPromptInfo.upscaler_raw = promptInfo.hires_upscaler || ""; + richPromptInfo.upscaler_raw = promptInfo.hires_upscaler; richPromptInfo.upscaler = findUpscaler(promptInfo.hires_upscaler); const { model, version } = findModelByHash(promptInfo.model_hash) richPromptInfo.model_used = model; richPromptInfo.model_used_version = version; + richPromptInfo.image_hash = await generateSha256Hash(arrayBuffer) + return richPromptInfo; } -export default parseMetadata; \ No newline at end of file +export default parseMetadata;