// ==UserScript==
// @name         NovelAI Autosave to Local Server
// @match        https://novelai.net/*
// @match        https://*.novelai.net/*
// @grant        GM_xmlhttpRequest
// @grant        GM_getValue
// @grant        GM_setValue
// @grant        GM_registerMenuCommand
// @connect      127.0.0.1
// @run-at       document-start
// @require      https://cdn.jsdelivr.net/npm/[email protected]/umd/index.js
// ==/UserScript==

(function () {
  'use strict';
  const uw = typeof unsafeWindow !== 'undefined' ? unsafeWindow : window;
  let PORT = GM_getValue('nai_autosave_port', 8891);
  const ENDPOINT = () => `http://127.0.0.1:${PORT}/`;
  const TARGET = '/ai/generate-image';

  function checkConnection() {
    GM_xmlhttpRequest({
      method: 'GET',
      url: ENDPOINT(),
      timeout: 4000,
      onload: (resp) => {
        if (resp.status === 200) {
          showConnectedToast();
        } else {
          warnUnreachable();
        }
      },
      onerror: warnUnreachable,
      ontimeout: warnUnreachable,
    });
  }

  function showConnectedToast() {
    const toast = document.createElement('div');
    toast.textContent = `✓ NovelAI autosave connected (port ${PORT})`;
    Object.assign(toast.style, {
      position: 'fixed', top: '10px', right: '10px', zIndex: 999999,
      background: '#27ae60', color: '#fff', padding: '8px 14px',
      borderRadius: '6px', fontFamily: 'sans-serif', fontSize: '12px',
      boxShadow: '0 2px 8px rgba(0,0,0,.4)',
    });
    document.body.appendChild(toast);
    setTimeout(() => toast.remove(), 3000);
  }

  GM_registerMenuCommand('NovelAI Autosave: set server port', () => {
    const input = prompt('Local autosave server port:', PORT);
    if (input && /^\d+$/.test(input.trim())) {
      PORT = parseInt(input.trim(), 10);
      GM_setValue('nai_autosave_port', PORT);
      checkConnection();
    }
  });

  let lastWarnTime = 0;
  function warnUnreachable() {
    const now = Date.now();
    if (now - lastWarnTime < 10000) return; // one banner per 10s max
    lastWarnTime = now;
    const existing = document.getElementById('nai-autosave-warning');
    if (existing) existing.remove();
    const banner = document.createElement('div');
    banner.id = 'nai-autosave-warning';
    banner.textContent = `⚠ NovelAI autosave: can't reach local server on port ${PORT}. Is it running? (click to dismiss)`;
    Object.assign(banner.style, {
      position: 'fixed', top: '10px', right: '10px', zIndex: 999999,
      background: '#c0392b', color: '#fff', padding: '10px 16px',
      borderRadius: '6px', fontFamily: 'sans-serif', fontSize: '13px',
      boxShadow: '0 2px 8px rgba(0,0,0,.4)', cursor: 'pointer', maxWidth: '300px',
    });
    banner.addEventListener('click', () => banner.remove());
    document.body.appendChild(banner);
    setTimeout(() => banner.remove(), 15000);
  }

  window.addEventListener('DOMContentLoaded', checkConnection);

  function toBase64(bytes) {
    let binary = '';
    const chunk = 0x8000;
    for (let i = 0; i < bytes.length; i += chunk) {
      binary += String.fromCharCode.apply(null, bytes.subarray(i, i + chunk));
    }
    return btoa(binary);
  }

function detectFormat(bytes) {
  if (bytes.length >= 8 && bytes[0] === 0x89 && bytes[1] === 0x50 && bytes[2] === 0x4E && bytes[3] === 0x47) {
    return 'png';
  }
  if (bytes.length >= 12 && bytes[0] === 0x52 && bytes[1] === 0x49 && bytes[2] === 0x46 && bytes[3] === 0x46 &&
      bytes[8] === 0x57 && bytes[9] === 0x45 && bytes[10] === 0x42 && bytes[11] === 0x50) {
    return 'webp';
  }
  return 'unknown';
}

  function readPngTextChunks(bytes) {
    const chunks = {};
    let offset = 8;
    while (offset < bytes.length) {
      const view = new DataView(bytes.buffer, bytes.byteOffset + offset, 8);
      const length = view.getUint32(0);
      const type = String.fromCharCode(bytes[offset+4], bytes[offset+5], bytes[offset+6], bytes[offset+7]);
      const dataStart = offset + 8, dataEnd = dataStart + length;
      if (type === 'tEXt') {
        const d = bytes.subarray(dataStart, dataEnd);
        const nullIdx = d.indexOf(0);
        chunks[new TextDecoder('latin1').decode(d.subarray(0, nullIdx))] =
          new TextDecoder('latin1').decode(d.subarray(nullIdx + 1));
      } else if (type === 'iTXt') {
        const d = bytes.subarray(dataStart, dataEnd);
        const kwEnd = d.indexOf(0);
        const keyword = new TextDecoder('utf-8').decode(d.subarray(0, kwEnd));
        const compressed = d[kwEnd + 1] === 1;
        let pos = d.indexOf(0, kwEnd + 3) + 1;
        pos = d.indexOf(0, pos) + 1;
        let textBytes = d.subarray(pos);
        if (compressed) textBytes = fflate.unzlibSync(textBytes);
        chunks[keyword] = new TextDecoder('utf-8').decode(textBytes);
      } else if (type === 'IDAT') {
        break;
      }
      offset = dataEnd + 4;
    }
    return chunks;
  }

  const isZip = (b) => b.length > 4 && b[0] === 0x50 && b[1] === 0x4B;
  const isPng = (b) => b.length > 8 && b[0] === 0x89 && b[1] === 0x50 && b[2] === 0x4E && b[3] === 0x47;
  const PNG_SIG = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];

  function findPngStarts(bytes) {
    const starts = [];
    outer:
    for (let i = 0; i <= bytes.length - 8; i++) {
      if (bytes[i] !== PNG_SIG[0]) continue;
      for (let j = 1; j < 8; j++) {
        if (bytes[i + j] !== PNG_SIG[j]) continue outer;
      }
      starts.push(i);
    }
    return starts;
  }

  function extractPngAt(bytes, start) {
    let offset = start + 8;
    while (offset + 8 <= bytes.length) {
      const view = new DataView(bytes.buffer, bytes.byteOffset + offset, 8);
      const length = view.getUint32(0);
      const type = String.fromCharCode(bytes[offset+4], bytes[offset+5], bytes[offset+6], bytes[offset+7]);
      const chunkEnd = offset + 8 + length + 4; // data + CRC
      if (chunkEnd > bytes.length) return null; // truncated, bail
      if (type === 'IEND') return bytes.subarray(start, chunkEnd);
      offset = chunkEnd;
    }
    return null; // no IEND found
  }

  function findRiffStarts(bytes) {
    const starts = [];
    for (let i = 0; i <= bytes.length - 12; i++) {
      if (bytes[i] !== 0x52 || bytes[i+1] !== 0x49 || bytes[i+2] !== 0x46 || bytes[i+3] !== 0x46) continue;
      if (bytes[i+8] !== 0x57 || bytes[i+9] !== 0x45 || bytes[i+10] !== 0x42 || bytes[i+11] !== 0x50) continue;
      starts.push(i);
    }
    return starts;
  }

  function extractWebpAt(bytes, start) {
    if (start + 8 > bytes.length) return null;
    const view = new DataView(bytes.buffer, bytes.byteOffset + start + 4, 4);
    const end = start + 8 + view.getUint32(0, true); // size field is little-endian in RIFF
    if (end > bytes.length) return null; // truncated
    return bytes.subarray(start, end);
  }

  function extractCompletePngs(bytes) {
    const out = [];
    for (const start of findPngStarts(bytes)) {
      const data = extractPngAt(bytes, start);
      if (data) out.push({ data, start, format: 'png' });
    }
    return out;
  }

  function extractCompleteWebps(bytes) {
    const out = [];
    for (const start of findRiffStarts(bytes)) {
      const data = extractWebpAt(bytes, start);
      if (data) out.push({ data, start, format: 'webp' });
    }
    return out;
  }

  function extractCompleteImages(bytes) {
    return [...extractCompletePngs(bytes), ...extractCompleteWebps(bytes)]
      .sort((a, b) => a.start - b.start);
  }

  async function handleGeneration(action, bytes, fallback) {
    let images;
    if (isZip(bytes)) {
      images = Object.values(fflate.unzipSync(bytes)).map((data) => ({ data, format: detectFormat(data) }));
    } else {
      images = extractCompleteImages(bytes);
    }
    if (images.length === 0) {
      console.warn('[NAI autosave] no image data found in response, first bytes:', Array.from(bytes.slice(0, 8)));
      return;
    }

    let finals = images
      .map((img) => ({ ...img, meta: img.format === 'png' ? readPngTextChunks(img.data) : {} }))
      .filter((p) => p.meta.Comment); // WebP entries never have this — chunk reading is PNG-only

    let usedFallback = false;
    if (finals.length === 0 && fallback) {
      usedFallback = true;
      finals = images.slice(-fallback.nSamples).map((img) => ({ ...img, meta: fallback.chunks }));
    }

    if (finals.length === 0) {
      console.warn('[NAI autosave] found image(s) but none looked like a finished image');
      return;
    }

    const subfolder = action === 'generate' ? 'txt2img-images' : 'img2img-images';
    for (const { data, meta, format } of finals) {
      GM_xmlhttpRequest({
        method: 'POST',
        url: ENDPOINT(),
        headers: { 'Content-Type': 'application/json' },
        data: JSON.stringify({ data: toBase64(data), metadata: meta, subfolder, ext: format }),
        timeout: 5000,
        onerror: warnUnreachable,
        ontimeout: warnUnreachable,
        onload: (resp) => { if (resp.status !== 200) { console.warn('[NAI autosave] server error', resp.status, resp.responseText); warnUnreachable(); } },
      });
    }
    console.log(`[NAI autosave] found ${images.length} embedded image(s) [${images.map(i => i.format).join(',')}], saved ${finals.length}${usedFallback ? ' (fallback metadata, no embedded chunk)' : ''}, action=${action}`);
  }

  function buildFallbackMeta(reqJson) {
    const p = reqJson.parameters || {};
    const comment = {
      prompt: reqJson.input || '',
      uc: p.negative_prompt || '',
      steps: p.steps,
      sampler: p.sampler,
      noise_schedule: p.noise_schedule,
      scale: p.scale,
      seed: p.seed,
      width: p.width,
      height: p.height,
      n_samples: p.n_samples || 1,
      model_name: reqJson.model || '', // raw internal id (e.g. "nai-diffusion-5-full") — no display-name mapping available client-side
      v4_prompt: p.v4_prompt || null,
      v4_negative_prompt: p.v4_negative_prompt || null,
    };
    return { chunks: { Comment: JSON.stringify(comment), Description: comment.prompt }, nSamples: comment.n_samples };
  }

  async function extractRequestInfo(body) {
    let json = null;
    try {
      if (body instanceof FormData) {
        const field = body.get('request');
        if (field) json = JSON.parse(await field.text());
      } else if (typeof body === 'string') {
        json = JSON.parse(body);
      }
    } catch (e) { /* leave json null, defaults below apply */ }

    return {
      action: json?.action || 'generate',
      fallback: json ? buildFallbackMeta(json) : null,
    };
  }

  // --- fetch hook ---
  const origFetch = uw.fetch;
  uw.fetch = async function (...args) {
    const resp = await origFetch.apply(this, args);
    const url = typeof args[0] === 'string' ? args[0] : args[0]?.url || '';
    if (url.includes(TARGET)) {
      console.log('[NAI autosave] caught fetch to', url);
      const { action, fallback } = await extractRequestInfo(args[1]?.body);
      const buf = await resp.clone().arrayBuffer();
      handleGeneration(action, new Uint8Array(buf), fallback);
    }
    return resp;
  };

  // --- XHR hook (in case the page uses XHR for upload-progress tracking) ---
  const origOpen = uw.XMLHttpRequest.prototype.open;
  const origSend = uw.XMLHttpRequest.prototype.send;
  uw.XMLHttpRequest.prototype.open = function (method, url, ...rest) {
    this.__nai_url = url;
    return origOpen.call(this, method, url, ...rest);
  };
  uw.XMLHttpRequest.prototype.send = function (body) {
    if (typeof this.__nai_url === 'string' && this.__nai_url.includes(TARGET)) {
      console.log('[NAI autosave] caught XHR to', this.__nai_url);
      const actionPromise = extractRequestInfo(body);
      this.addEventListener('load', async () => {
        const { action, fallback } = await actionPromise;
        let bytes;
        if (this.responseType === 'blob') {
          bytes = new Uint8Array(await this.response.arrayBuffer());
        } else if (this.responseType === 'arraybuffer') {
          bytes = new Uint8Array(this.response);
        } else {
          console.warn('[NAI autosave] XHR responseType is', this.responseType, '- cannot read binary safely');
          return;
        }
        handleGeneration(action, bytes, fallback);
      });
    }
    return origSend.call(this, body);
  };
})();
Edit

Pub: 21 Aug 2026 12:15 UTC

Edit: 23 Aug 2026 15:21 UTC

Views: 12