☰

ExportCivitAIMetadata

导出 civitai.com 的 safetensors 模型元数据 / Export .safetensor file's metadata from civitAI

当前为 2023-06-11 提交的版本，查看 最新版本。

// ==UserScript==
// @name         ExportCivitAIMetadata
// @namespace    https://github.com/magicFeirl/ExportCivitAIMetadata.git
// @description  导出 civitai.com 的 safetensors 模型元数据 / Export .safetensor file's metadata from civitAI
// @author       ctrn43062
// @match        https://civitai.com/*
// @icon         https://www.google.com/s2/favicons?sz=64&domain=civitai.com
// @version      0.2
// @note         0.2 修复某些情况下复制按钮没有出现的bug
// @note         0.1 init
// @license      MIT
// ==/UserScript==

/**
* Usage:
* SafetensorsHeaderReader.readFromURL(url, [offset])
**/
class SafetensorsHeaderReader {
  static async getHeaderLengthFromURL(url, offset = 0) {
    const resp = await fetch(url, { headers: { range: 'bytes=0-7' } })

    try {
      const buffer = new BigUint64Array(await resp.arrayBuffer()).buffer
      return new DataView(buffer).getBigUint64(0, true) + offset
    } catch (error) {
      // backend error?
      console.log('Failed to get header length from backend:', error)
      return -1
    }

  }

  static async readFromURL(url, offset = 0) {
    if (!url) {
      return console.error('No model url')
    }

    const headerLength = await this.getHeaderLengthFromURL(url, BigInt(offset))

    if (headerLength <= 0) {
      return {
        error: 'Failed to get header length from civitai'
      }
    }

    const resp = await fetch(url, { headers: { range: `bytes=8-${headerLength}` } })
    const data = await resp.json()
    return data
  }
  // TODO: readHeader from file
}

function createCopyHeaderButton() {
  const icon = `<?xml version="1.0" ?><!-- Uploaded to: SVG Repo, www.svgrepo.com, Generator: SVG Repo Mixer Tools -->
    <svg width="800px" height="800px" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M10 8V7C10 6.05719 10 5.58579 10.2929 5.29289C10.5858 5 11.0572 5 12 5H17C17.9428 5 18.4142 5 18.7071 5.29289C19 5.58579 19 6.05719 19 7V12C19 12.9428 19 13.4142 18.7071 13.7071C18.4142 14 17.9428 14 17 14H16M7 19H12C12.9428 19 13.4142 19 13.7071 18.7071C14 18.4142 14 17.9428 14 17V12C14 11.0572 14 10.5858 13.7071 10.2929C13.4142 10 12.9428 10 12 10H7C6.05719 10 5.58579 10 5.29289 10.2929C5 10.5858 5 11.0572 5 12V17C5 17.9428 5 18.4142 5.29289 18.7071C5.58579 19 6.05719 19 7 19Z" stroke="#fff" stroke-linecap="round" stroke-linejoin="round"/></svg>`

  const downloadBtn = document.querySelector('.mantine-Group-root .mantine-Stack-root a[download]')

  if (!downloadBtn) {
    // waiting for DOM loaded
    console.error('can\'t find download button')
    return {}
  }

  const buttonWrapper = downloadBtn.parentElement.parentElement

  const btnCopy = buttonWrapper.lastChild.cloneNode(true)

  btnCopy.setAttribute('id', 'CVI-copy-btn')
  btnCopy.setAttribute('title', 'Export model\'s metadata')

  const svg = btnCopy.querySelector('span.mantine-Button-label')
  svg.innerHTML = icon

  buttonWrapper.appendChild(btnCopy)

  return { header: buttonWrapper.parentElement, btn: btnCopy.firstChild }
}

function download(url, name) {
  const link = document.createElement('a');
  link.href = url;
  link.download = name;
  document.body.appendChild(link);
  link.click();
  setTimeout(() => link.remove(), 0)
}

function downloadText(text, name) {
  const blob = new Blob([text])
  const url = URL.createObjectURL(blob)
  download(url, name)
  URL.revokeObjectURL(url)
}

function init() {
  // we have init the copy button
  if (document.querySelector('#CVI-copy-btn')) {
    return;
  }

  const { header = undefined, btn = undefined } = createCopyHeaderButton()

  if (!header || !btn) {
    return;
  }

  btn.onclick = () => {
    const isSafetensor = [...header.querySelectorAll('.mantine-Text-root')].map(el => el.textContent).some(text => text.toLowerCase() === 'safetensor')

    function stringifyObject(obj) {
      return JSON.stringify(obj, (_, v) => {
        if (typeof v === 'string') {
          try {
            return JSON.parse(v, null, 2)
          } catch {
            return v
          }
        }

        return v
      }, 2)
    }

    if (!isSafetensor) {
      alert('Not a .safetensors model')
      return;
    }

    const modelURL = header.querySelector('a[download]').href
    const RANGE_OFFSET = 7

    const { pathname } = location
    // model title + model url id = filename
    const title = document.querySelector('#freezeBlock h1').textContent
    const path = pathname.substring(pathname.lastIndexOf('/') + 1)
    const filename = title + '_' + path + '_metadata.json'

    btn.setAttribute('disabled', true)
    SafetensorsHeaderReader.readFromURL(modelURL, RANGE_OFFSET).then((json) => {
      if (json['error']) {
        alert('error: ' + json['error'] + ', please try later.')
        return;
      }

      return json
    }).catch((error) => {
      alert('error: Network error:', error)
      console.error(error)
    }).then(json => {
      if (!json) {
        return
      }

      // 避免 undefined 或 null
      json.__metadata__ = json.__metadata__ || {}

      if (!Object.keys(json.__metadata__).length) {
        alert('This model has no metadata')
        return
      }
      // export __metadata__ to .txt file
      downloadText(stringifyObject(json.__metadata__), filename)
    }).finally(() => {
      btn.removeAttribute('disabled')
    })
  }
}

function initPageChangeObserver(callback, window) {
  const ob = new window.MutationObserver((mutationList) => {
    if (mutationList.some(record => record.type === 'childList')) {
      // model detail page
      if (/civitai\.com\/models\/\d+/.test(location.href)) {
        callback && callback()
      }
    }
  })

  ob.observe(window.document.querySelector('title'), { childList: true, subtree: true })

  return ob
}

function findCopyButton(callback) {
  setTimeout(() => {
    if (!document.querySelector('#CVI-copy-btn')) {
      findCopyButton(callback)
      callback()
    }
  }, 50)
}

(function () {
  findCopyButton(init)

  // For SPA
  initPageChangeObserver(init, unsafeWindow || window)
})();