NovelAI / Stable Diffusion PNG Chunk Reader

chunk_reader.js

module.exports = chunk_reader
const extract = require('png-chunks-extract')

function DiffusionObj (positive, negative, software, sampler, seed, steps, scale, strength, noise) {
  this.positive = positive
  this.negative = negative
  this.software = software
  this.sampler = sampler
  this.seed = seed
  this.steps = steps
  this.scale = scale
  this.strength = strength
  this.noise = noise
}

function chunk_reader (data) {
  let diffusion_obj
  const chunks = extract(new Uint8Array(data))
  const selected_chunks = selectChunks(chunks, ["tEXt", "iTXt"])
  const decoded_chunks = decodeChunks(selected_chunks)
  if (isNovelAI(decoded_chunks)) {
    return parseNovelAI(decoded_chunks)
  } else if (isStableDiffusion(decoded_chunks)) {
    return parseStableDiffusion(decoded_chunks)
  } else {
    return new DiffusionObj("NaN", "NaN", "NaN", "NaN", "NaN", "NaN", "NaN", "NaN", "NaN")
  }
}

function selectChunks (chunks, select_list) {
  let selected_chunks = []
  chunks.forEach(chunk => {
    select_list.forEach(name => {
      if (chunk.name === name) {
        selected_chunks.push(chunk)
      }
    })    
  })
  return selected_chunks
}

function decodeChunks (chunks) {
  return chunks.map(data => {
    if (data.data && data.name) {
      data = data.data
    }
  
    var naming = true
    var text_candidate = []
    var name = ''
  
    for (var i = 0; i < data.length; i++) {
      var code = data[i]
  
      if (naming) {
        if (code) {
          name += String.fromCharCode(code)
        } else {
          naming = false
        }
      } else {
        if (code) {
          text_candidate.push(code)
        } else {
          // do nothing
        }
      }
    }
  
    return {
      keyword: name,
      text: new TextDecoder("utf-8").decode(Uint8Array.from(text_candidate).buffer)
    }
  })
}

function isNovelAI (chunks) {
  let is_novelai = false
  chunks.forEach(chunk => {
    if (chunk.text == "NovelAI") {
      is_novelai = true
    }
  })
  return is_novelai
}

function isStableDiffusion (chunks) {
  let is_stablediffusion = false
  chunks.forEach(chunk => {
    if (chunk.keyword == "parameters") {
      is_stablediffusion = true
    }
  })
  return is_stablediffusion
}

function parseNovelAI (chunks) {
  let diff = new DiffusionObj()
  const params = JSON.parse(chunks[4].text);
  diff.positive = chunks[1].text
  diff.negative = params.uc
  diff.software = chunks[2].text
  diff.sampler = params.sampler
  diff.seed = params.seed
  diff.steps = params.steps
  diff.scale = params.scale
  diff.strength = params.strength
  diff.noise = params.noise
  return diff
}

function parseStableDiffusion (chunks) {
  let diff = new DiffusionObj()
  const data = chunks[0].text.split(/\n/)
  diff.positive = data[0]
  diff.negative = data[1].split("Negative prompt: ")[1]
  other_data = data[2].split(", ")
  diff.software = "Stable Diffusion"
  diff.sampler = other_data[1].split("Sampler: ")[1]
  diff.seed = other_data[3].split("Seed: ")[1]
  diff.steps = other_data[0].split("Steps: ")[1]
  diff.scale = other_data[2].split("CFG scale: ")[1]
  diff.strength = "-"
  diff.noise = "-"
  return diff
}