☰

Prompt Extractor

이미지로부터 Exif 정보를 추출해 사용자에게 보여줍니다

// ==UserScript==
// @name        Prompt Extractor
// @namespace   https://github.com/toriato/userscripts/prompt-extractor.user.js
// @version     0.1.3
// @description 이미지로부터 Exif 정보를 추출해 사용자에게 보여줍니다
// @author      Sangha Lee <[email protected]>
// @license     MIT
// @match       https://arca.live/b/*
// @run-at      document-start
// @grant       GM_xmlhttpRequest
// @grant       GM_addStyle
// @require     https://unpkg.com/exifreader/dist/exif-reader.js
// ==/UserScript==

GM_addStyle(/*css*/`
  @keyframes spin {
    from { transform:rotate(0deg) }
    to { transform:rotate(360deg) }
  }

  figure.params {
    margin: 0;
    position: relative;
    display: table;
  }
  figure.params img {
    max-width: 100%;
  }

  /* 우측 상단 상태 아이콘 */
  figure.params:not([data-params=""])::after {
    position: absolute;
    right: 0;
    top: 0;
    margin: .5em;
    font-size: 2rem;
    text-shadow: 0 0 4px black;
    content: '❤️'
  }
  figure.params.loading::after {
    animation: spin 1s infinite linear;
    content: '🌀'
  }
  figure.params:not(.loading):not([data-params=""]):hover::after {
    display: none;
  }

  figure.params figcaption {
    transition: transform .25s, opacity .25s;
    transform: scaleY(0);
    transform-origin: top;
    position: absolute;
    left: 0;
    top: 0;
    overflow-y: auto;
    max-height: 50%;
    padding: .5em;
    opacity: 0;
    background-color: rgba(0, 0, 0, 0.5);
    text-align: left;
    pointer-events: none;
  }
  figure.params:not(.loading):not([data-params=""]):hover figcaption {
    transform: scaleY(1);
    opacity: 1;
    pointer-events: inherit;
  }
`)

/**
 * UPNG.js - JS PNG Decoder/Encoder
 * https://github.com/photopea/UPNG.js
 * MIT License
 */
class UPNG {
  static bin = {
    nextZero: (data, p) => {
      while (data[p] != 0) p++
      return p
    },
    readUshort: (buff, p) =>
      (buff[p] << 8) | buff[p + 1],
    writeUshort: (buff, p, n) => {
      buff[p] = (n >> 8) & 255
      buff[p + 1] = n & 255
    },
    readUint: (buff, p) =>
      (buff[p] * (256 * 256 * 256)) + ((buff[p + 1] << 16) | (buff[p + 2] << 8) | buff[p + 3]),
    writeUint: (buff, p, n) => {
      buff[p] = (n >> 24) & 255
      buff[p + 1] = (n >> 16) & 255
      buff[p + 2] = (n >> 8) & 255
      buff[p + 3] = n & 255
    },
    readASCII: (buff, p, l) => {
      let s = ''
      for (let i = 0; i < l; i++)
        s += String.fromCharCode(buff[p + i])
      return s
    },
    writeASCII: (data, p, s) => {
      for (let i = 0; i < s.length; i++)
        data[p + i] = s.charCodeAt(i)
    },
    readBytes: (buff, p, l) => {
      const arr = []
      for (let i = 0; i < l; i++)
        arr.push(buff[p + i])
      return arr
    },
    pad: (n) =>
      n.length < 2 ? '0' + n : n,
    readUTF8: function (buff, p, l) {
      let s = ''
      let ns
      for (var i = 0; i < l; i++) s += '%' + UPNB.bin.pad(buff[p + i].toString(16))
      try { ns = decodeURIComponent(s) }
      catch (e) { return UPNG.bin.readASCII(buff, p, l) }
      return ns
    }
  }

  static magicNumbers = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]

  static decode(buff) {
    const bin = UPNG.bin
    const data = new Uint8Array(buff)
    const texts = {}

    for (let i = 0; i < 8; i++) {
      if (data[i] !== UPNG.magicNumbers[i]) {
        throw Error('Input file is not a PNG')
      }
    }

    let offset = 8
    while (offset < data.length) {
      const len = bin.readUint(data, offset); offset += 4;
      const type = bin.readASCII(data, offset, 4); offset += 4;

      // 스펙 상 tEXt 청크는 순서 관계 없으나 프롬프트는 상단에 위치하므로
      // 빠른 처리를 위해 데이터가 시작되면 중단함
      // https://www.w3.org/TR/2003/REC-PNG-20031110/#5ChunkOrdering
      if (type === 'IDAT') {
        break
      }

      if (type === 'tEXt') {
        const nz = bin.nextZero(data, offset);
        const keyword = bin.readASCII(data, offset, nz - offset);
        const textLen = offset + len - nz - 1;
        texts[keyword] = bin.readASCII(data, nz + 1, textLen)
      }

      offset += len + 4;
    }

    return texts
  }
}

/**
 * 파라미터 문자열 파싱에 사용되는 정규표현식 패턴
 * 
 * https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/0cc0ee1bcb4c24a8c9715f66cede06601bfc00c8/modules/generation_parameters_copypaste.py#L15
 */
const paramsPattern = /\s*([\w ]+):\s*("(?:\\"[^,]|\\"|\\|[^\"])+"|[^,]*)(?:,|$)/g

/**
 * 생성에 사용된 파라미터 문자열을 키, 값 형식의 Object 로 파싱합니다.
 * 
 * https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/0cc0ee1bcb4c24a8c9715f66cede06601bfc00c8/modules/generation_parameters_copypaste.py#LL226C13-L226C13
 * @param {string} str
 * @returns {Object.<string, string>}
 */
function parseGenerationParams(str) {
  const lines = str.trim().split('\n')
  const res = Object.fromEntries(
    // 반환 값: [ 전체 문자열, 키, 값 ]
    [...lines.pop().matchAll(paramsPattern)]
      // 첫번째 값은 일치한 전체 문자열이므로 필요 없음
      .map(v => v.slice(1))
  )

  // 프롬프트와 부정 프롬프트는 둘 다 여러 줄일 수 있기 때문에 반복문으로 확인해야 함
  let key = 'Prompt'

  for (let line of lines) {
    // 맨 앞 문자열이 일치하면 그 때부터 네거티브 프롬프트로 처리하는데...
    // 일반 프롬프트에 동일한 문자열이 존재하면 오작동하지 않을까?
    // 근데 자동좌 레포지토리에서도 이렇게 처리하니까 아무튼 내 잘못 아님
    // https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/0cc0ee1bcb4c24a8c9715f66cede06601bfc00c8/modules/generation_parameters_copypaste.py#L251
    if (line.startsWith('Negative prompt:')) {
      key = 'Negative Prompt'
      line = line.slice(16).trim()
    }

    // 없으면 새 문자열 만들고 있으면 새 줄 넣고 추가하기
    if (key in res) {
      res[key] += '\n' + line
    } else {
      res[key] = line
    }
  }

  return res
}

/**
 * 이미지가 모두 불러와졌을 때 실행되는 이벤트 함수입니다.
 * 
 * @param {UIEvent} event
 */
function onLoad(event) {
  /** @type {HTMLImageElement} */
  const node = event.target

  // 작은 이미지는 메타데이터 확인하지 않기
  const rect = node.getBoundingClientRect()
  if (rect.width < 128 || rect.height < 128) {
    return
  }

  let src = new URL(node.src)

  // 아카라이브에선 원본 이미지에만 Exif 데이터가 존재함 
  if (src.host.endsWith('namu.la')) {
    src.searchParams.set('type', 'orig')
  }

  // 기존 이미지 요소 위에 파라미터를 표시하기 위해 figure 요소로 감싸기
  const $figure = document.createElement('figure')
  $figure.classList.add('params', 'loading')
  $figure.innerHTML = /*html*/`
    ${node.closest('p').innerHTML}
    <figcaption></figcaption>
  `

  node.closest('p').replaceWith($figure)

  // Exif 로부터 파라미터 문자열 가져오기
  let params = ''

  // 이미 불러온 이미지로는 데이터를 가져올 수 없기 때문에 새 요청을 만들 필요가 있음
  // 브라우저가 캐시해줄테니 속도에 큰 지장을 주진 않을거임... 아마도...?
  new Promise((resolve, reject) => {
    GM_xmlhttpRequest({
      url: src.toString(),
      responseType: 'arraybuffer',
      onload: resolve,
      onerror: reject
    })
  })
    .then(res => {
      const headers = Object.fromEntries(
        res.responseHeaders
          .split(/\r?\n/)
          .map(v => {
            const [key, value] = v.split(':', 2).map(v => v.trim())
            return [key.toLowerCase(), value]
          })
      )

      const contentType = headers['content-type']
      switch (contentType) {
        // PNG 는 Exif 가 아닌 tEXt 키워드를 통해 파라미터가 저장되기 때문에
        // UPNG 라이브러리를 통해 파라미터 문자열을 가져올 수 있음
        case 'image/png':
          const texts = UPNG.decode(res.response)
          params = texts['parameters'] ?? ''
          break

        // ExifReader 라이브러리를 사용해 Exif 중 UserComment 로부터 파라미터 가져오기
        // https://github.com/mattiasw/ExifReader
        default:
          // 반환 받은 파일이 이미지가 아니라면 무시하기
          if (!contentType.startsWith('image/')) {
            return
          }

          try {
            const tags = ExifReader.load(res.response)
            if (tags?.UserComment?.value) {
              params = String.fromCharCode(
                // 첫 8바이트는 인코딩 타입이므로 디코딩 할 필요 없음
                // https://www.awaresystems.be/imaging/tiff/tifftags/privateifd/exif/usercomment.html
                ...tags.UserComment.value.slice(8).filter(v => v !== 0)
              )
            }

          } catch (err) {
            // 메타데이터가 존재하지 않는 이미지라면 무시하기
            if (err.name !== 'MetadataMissingError') {
              throw err
            }
          }
      }
    })

    // TODO: 깔끔한 오류 핸들링
    // .catch(err => ...)

    // figcaption 요소를 통해 파라미터 정보 표시하기
    .finally(() => {
      $figure.classList.remove('loading')
      $figure.dataset.params = params

      // 파라미터 값이 존재하지 않는다면 하위 요소 생성하지 않기
      if (!params) {
        return
      }

      // TODO: 파싱한 파라미터 표로 보여주고 복사하는 기능 만들기
      // const parsedParams = parseGenerationParams(params)
      // console.log(parsedParams)

      $figure.querySelector('figcaption').innerHTML = params
    })
}

// 새로 추가되는 이미지 요소에 load 이벤트 등록하기
new MutationObserver(mutations => {
  for (let mutation of mutations) {
    for (let node of mutation.addedNodes) {
      // 노드가 이미지 태그가 아니라면 무시하기
      if (!(node instanceof HTMLImageElement)) {
        continue
      }

      node.addEventListener('load', onLoad)
    }
  }
}).observe(
  document,
  {
    attributes: true,
    childList: true,
    subtree: true
  }
)