NovelAI image gen request script

nai.mjs

// Requires node-fetch. Just do `npm install node-fetch`
import fetch from 'node-fetch';
import fs from 'fs';


// This expects a 'key.txt' file with your nai api key in it. You can get that by peeking at a network request
const API_KEY = fs.readFileSync("./key.txt").toString();
const API_HEADER = "Bearer " + API_KEY;

// This is the default preset I think? I forget if I modified this slightly
const UNDESIRED_CONTENT = "nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry";

// The full model
const DEFAULT_MODEL = "nai-diffusion";
const FURRY_MODEL = "nai-diffusion-furry";

function get_image_params(name, seed, UC, scale) {
    if (!IMAGE_MODES.hasOwnProperty(name)) {
        throw new Error("Invalid mode name");
    }
    let mode = IMAGE_MODES[name];

    // Merge them
    return Object.assign({
        seed,
        // 1 sample is free
        n_samples: 1,
        // Typical
        sampler: "k_euler_ancestral",
        scale,
        steps: 28,
        // These only appear in enhance for the UI but they seemed to be sent?    
        noise: 0.2,
        strength: 0.7,

        uc: UC,
        // They send this, but it seems that they always include the UC preset manually?
        ucPreset: 0,
    }, mode)
}
const IMAGE_MODES = {
    PortraitNormal: {
        height: 768,
        width: 512,
    },
    LandscapeNormal: {
        height: 512,
        width: 768
    }
};

async function request_image(input, mode, seed, model = DEFAULT_MODEL) {
    const IMAGE_API_URL = "https://api.novelai.net/ai/generate-image";
    // They temp used this link during the outages
    //const IMAGE_API_URL = "https://backend-production-svc.novelai.net/ai/generate-image";

    // Add the common prompt prefix
    input = "masterpiece, best quality" + input;

    let scale;
    let uc;
    // I found that scale 8 and uc="" works better for furry, so I just set it here
    if (model == FURRY_MODEL) {
        scale = 8;
        uc = "";
    } else {
        scale = 11;
        uc = UNDESIRED_CONTENT;
    }

    const body = JSON.stringify({
        input,
        model,
        parameters: get_image_params(mode, seed, uc, scale)
    });

    // TODO: We could do better than this if there was nice impl of event sources or something
    // for nodejs...

    let resp;
    let data;

    resp = await fetch(IMAGE_API_URL, {
        method: "POST",
        body,
        headers: {
            'Authorization': API_HEADER,
            'Content-Type': 'application/json',
            'accept': 'application/json',
        },
    });
    data = await resp.text();

    // This is a hacky way of not getting a library or implementing a custom parser for it
    // (There seems to not be that many? unless i'm misunderstanding the protocl)
    // This probably wouldn't work if it was generating multiple images at once
    if (data.startsWith("event: newImage")) {
        // Strip the prefix off
        let data1 = data.slice("event: newImage\n".length);
        let data2 = data1.slice(data1.indexOf("\n") + 1 + "data:".length);
        // Convert the base 64 to a buffer
        let data3 = Buffer.from(data2, "base64");
        return data3;
    } else {
        return null;
    }
}

async function request_image_to_file(filename, input, mode, seed, model = DEFAULT_MODEL) {
    let img = await request_image(input, mode, seed, model);
    if (img) {
        await fs.promises.writeFile(filename, img);
    } else {
        console.log("Failed to generate image");
    }
}


function generate_seed() {
    // This is how they do it in nai, last I checked
    return Math.floor(Math.random() * Math.pow(2, 32) - 1);
}

// Example
request_image_to_file("output.png", "toaster, rgb, gaming, neon, realistic", "PortraitNormal", generate_seed(), DEFAULT_MODEL);