Add NAI Diffusion upscaling. Add Anlas guard and view Anlas button

public/script.js

@@ -98,8 +98,8 @@ import {
     loadNovelPreset,
     loadNovelSettings,
     nai_settings,
-    setNovelData,
     adjustNovelInstructionPrompt,
+    loadNovelSubscriptionData,
 } from "./scripts/nai-settings.js";
 
 import {
@@ -635,7 +635,7 @@ let online_status = "no_connection";
 let api_server = "";
 let api_server_textgenerationwebui = "";
 //var interval_timer = setInterval(getStatus, 2000);
-let interval_timer_novel = setInterval(getStatusNovel, 90000);
+//let interval_timer_novel = setInterval(getStatusNovel, 90000);
 let is_get_status = false;
 let is_get_status_novel = false;
 let is_api_button_press = false;
@@ -5343,32 +5343,19 @@ export async function displayPastChats() {
 //************************************************************
 async function getStatusNovel() {
     if (is_get_status_novel) {
-        const data = {};
+        try {
+            const result = await loadNovelSubscriptionData();
 
-        jQuery.ajax({
+            if (!result) {
-            type: "POST", //
+                throw new Error('Could not load subscription data');
-            url: "/getstatus_novelai", //
-            data: JSON.stringify(data),
-            beforeSend: function () {
-
-            },
-            cache: false,
-            dataType: "json",
-            contentType: "application/json",
-            success: function (data) {
-                if (data.error != true) {
-                    setNovelData(data);
-                    online_status = `${getNovelTier(data.tier)}`;
             }
-                resultCheckStatusNovel();
+
-            },
+            online_status = getNovelTier();
-            error: function (jqXHR, exception) {
+        } catch {
             online_status = "no_connection";
-                console.log(exception);
+        }
-                console.log(jqXHR);
+
         resultCheckStatusNovel();
-            },
-        });
     } else {
         if (is_get_status != true && is_get_status_openai != true) {
             online_status = "no_connection";

public/scripts/extensions/stable-diffusion/index.js

@@ -14,9 +14,10 @@ import {
 } from "../../../script.js";
 import { getApiUrl, getContext, extension_settings, doExtrasFetch, modules, renderExtensionTemplate } from "../../extensions.js";
 import { selected_group } from "../../group-chats.js";
-import { stringFormat, initScrollHeight, resetScrollHeight, timestampToMoment, getCharaFilename, saveBase64AsFile } from "../../utils.js";
+import { stringFormat, initScrollHeight, resetScrollHeight, getCharaFilename, saveBase64AsFile } from "../../utils.js";
 import { getMessageTimeStamp, humanizedDateTime } from "../../RossAscends-mods.js";
 import { SECRET_KEYS, secret_state } from "../../secrets.js";
+import { getNovelUnlimitedImageGeneration, getNovelAnlas, loadNovelSubscriptionData } from "../../nai-settings.js";
 export { MODULE_NAME };
 
 // Wraps a string into monospace font-face span
@@ -177,6 +178,13 @@ const defaultSettings = {
     hr_second_pass_steps_min: 0,
     hr_second_pass_steps_max: 150,
     hr_second_pass_steps_step: 1,
+
+    // NovelAI settings
+    novel_upscale_ratio_min: 1.0,
+    novel_upscale_ratio_max: 4.0,
+    novel_upscale_ratio_step: 0.1,
+    novel_upscale_ratio: 1.0,
+    novel_anlas_guard: false,
 }
 
 const getAutoRequestBody = () => ({ url: extension_settings.sd.auto_url, auth: extension_settings.sd.auto_auth });
@@ -226,6 +234,8 @@ async function loadSettings() {
     $('#sd_hr_scale').val(extension_settings.sd.hr_scale).trigger('input');
     $('#sd_denoising_strength').val(extension_settings.sd.denoising_strength).trigger('input');
     $('#sd_hr_second_pass_steps').val(extension_settings.sd.hr_second_pass_steps).trigger('input');
+    $('#sd_novel_upscale_ratio').val(extension_settings.sd.novel_upscale_ratio).trigger('input');
+    $('#sd_novel_anlas_guard').prop('checked', extension_settings.sd.novel_anlas_guard);
     $('#sd_horde').prop('checked', extension_settings.sd.horde);
     $('#sd_horde_nsfw').prop('checked', extension_settings.sd.horde_nsfw);
     $('#sd_horde_karras').prop('checked', extension_settings.sd.horde_karras);
@@ -389,6 +399,31 @@ async function onSourceChange() {
     await Promise.all([loadModels(), loadSamplers()]);
 }
 
+async function onViewAnlasClick() {
+    const result = await loadNovelSubscriptionData();
+
+    if (!result) {
+        toastr.warning('Are you subscribed?', 'Could not load NovelAI subscription data');
+        return;
+    }
+
+    const anlas = getNovelAnlas();
+    const unlimitedGeneration = getNovelUnlimitedImageGeneration();
+
+    toastr.info(`Free image generation: ${unlimitedGeneration ? 'Yes' : 'No'}`, `Anlas: ${anlas}`);
+}
+
+function onNovelUpscaleRatioInput() {
+    extension_settings.sd.novel_upscale_ratio = Number($('#sd_novel_upscale_ratio').val());
+    $('#sd_novel_upscale_ratio_value').text(extension_settings.sd.novel_upscale_ratio.toFixed(1));
+    saveSettingsDebounced();
+}
+
+function onNovelAnlasGuardInput() {
+    extension_settings.sd.novel_anlas_guard = !!$('#sd_novel_anlas_guard').prop('checked');
+    saveSettingsDebounced();
+}
+
 async function onHordeNsfwInput() {
     extension_settings.sd.horde_nsfw = !!$(this).prop('checked');
     saveSettingsDebounced();
@@ -642,7 +677,7 @@ async function loadAutoSamplers() {
 
 async function loadNovelSamplers() {
     if (!secret_state[SECRET_KEYS.NOVEL]) {
-        toastr.warning('NovelAI API key is not set.');
+        console.debug('NovelAI API key is not set.');
         return [];
     }
 
@@ -773,7 +808,7 @@ async function loadAutoModels() {
 
 async function loadNovelModels() {
     if (!secret_state[SECRET_KEYS.NOVEL]) {
-        toastr.warning('NovelAI API key is not set.');
+        console.debug('NovelAI API key is not set.');
         return [];
     }
 
@@ -1123,6 +1158,8 @@ async function generateAutoImage(prompt) {
  * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete.
  */
 async function generateNovelImage(prompt) {
+    const { steps, width, height } = getNovelParams();
+
     const result = await fetch('/api/novelai/generate-image', {
         method: 'POST',
         headers: getRequestHeaders(),
@@ -1130,11 +1167,12 @@ async function generateNovelImage(prompt) {
             prompt: prompt,
             model: extension_settings.sd.model,
             sampler: extension_settings.sd.sampler,
-            steps: extension_settings.sd.steps,
+            steps: steps,
             scale: extension_settings.sd.scale,
-            width: extension_settings.sd.width,
+            width: width,
-            height: extension_settings.sd.height,
+            height: height,
             negative_prompt: extension_settings.sd.negative_prompt,
+            upscale_ratio: extension_settings.sd.novel_upscale_ratio,
         }),
     });
 
@@ -1146,6 +1184,61 @@ async function generateNovelImage(prompt) {
     }
 }
 
+/**
+ * Adjusts extension parameters for NovelAI. Applies Anlas guard if needed.
+ * @returns {{steps: number, width: number, height: number}} - A tuple of parameters for NovelAI API.
+ */
+function getNovelParams() {
+    let steps = extension_settings.sd.steps;
+    let width = extension_settings.sd.width;
+    let height = extension_settings.sd.height;
+
+    // Don't apply Anlas guard if it's disabled.d
+    if (!extension_settings.sd.novel_anlas_guard) {
+        return { steps, width, height };
+    }
+
+    const MAX_STEPS = 28;
+    const MAX_PIXELS = 409600;
+
+    if (width * height > MAX_PIXELS) {
+        const ratio = Math.sqrt(MAX_PIXELS / (width * height));
+
+        // Calculate new width and height while maintaining aspect ratio.
+        var newWidth = Math.round(width * ratio);
+        var newHeight = Math.round(height * ratio);
+
+        // Ensure new dimensions are multiples of 64. If not, reduce accordingly.
+        if (newWidth % 64 !== 0) {
+            newWidth = newWidth - newWidth % 64;
+        }
+
+        if (newHeight % 64 !== 0) {
+            newHeight = newHeight - newHeight % 64;
+        }
+
+        // If total pixel count after rounding still exceeds MAX_PIXELS, decrease dimension size by 64 accordingly.
+        while (newWidth * newHeight > MAX_PIXELS) {
+            if (newWidth > newHeight) {
+                newWidth -= 64;
+            } else {
+                newHeight -= 64;
+            }
+        }
+
+        console.log(`Anlas Guard: Image size (${width}x${height}) > ${MAX_PIXELS}, reducing size to ${newWidth}x${newHeight}`);
+        width = newWidth;
+        height = newHeight;
+    }
+
+    if (steps > MAX_STEPS) {
+        console.log(`Anlas Guard: Steps (${steps}) > ${MAX_STEPS}, reducing steps to ${MAX_STEPS}`);
+        steps = MAX_STEPS;
+    }
+
+    return { steps, width, height };
+}
+
 async function sendMessage(prompt, image) {
     const context = getContext();
     const messageText = `[${context.name2} sends a picture that contains: ${prompt}]`;
@@ -1357,6 +1450,9 @@ jQuery(async () => {
     $('#sd_hr_scale').on('input', onHrScaleInput);
     $('#sd_denoising_strength').on('input', onDenoisingStrengthInput);
     $('#sd_hr_second_pass_steps').on('input', onHrSecondPassStepsInput);
+    $('#sd_novel_upscale_ratio').on('input', onNovelUpscaleRatioInput);
+    $('#sd_novel_anlas_guard').on('input', onNovelAnlasGuardInput);
+    $('#sd_novel_view_anlas').on('click', onViewAnlasClick);
     $('#sd_character_prompt_block').hide();
 
     $('.sd_settings .inline-drawer-toggle').on('click', function () {

public/scripts/extensions/stable-diffusion/settings.html

@@ -50,6 +50,18 @@
                 </label>
             </div>
             <div data-sd-source="novel">
+                <div class="flex-container">
+                    <label for="sd_novel_anlas_guard" class="checkbox_label flex1" title="Automatically adjust generation parameters to ensure free image generations.">
+                        <input id="sd_novel_anlas_guard" type="checkbox" />
+                        <span data-i18n="Avoid spending Anlas">
+                            Avoid spending Anlas
+                        </span>
+                        <span data-i18n="Opus tier" class="toggle-description">(Opus tier)</span>
+                    </label>
+                    <div id="sd_novel_view_anlas" class="menu_button menu_button_icon">
+                        View my Anlas
+                    </div>
+                </div>
                 <i>Hint: Save an API key in the NovelAI API settings to use it here.</i>
             </div>
             <label for="sd_scale">CFG Scale (<span id="sd_scale_value"></span>)</label>
@@ -84,6 +96,10 @@
                 <label for="sd_hr_second_pass_steps">Hires steps (2nd pass) (<span id="sd_hr_second_pass_steps_value"></span>)</label>
                 <input id="sd_hr_second_pass_steps" type="range" min="{{hr_second_pass_steps_min}}" max="{{hr_second_pass_steps_max}}" step="{{hr_second_pass_steps_max}}" value="{{hr_second_pass_steps}}" />
             </div>
+            <div data-sd-source="novel">
+                <label for="sd_novel_upscale_ratio">Upscale by (<span id="sd_novel_upscale_ratio_value"></span>)</label>
+                <input id="sd_novel_upscale_ratio" type="range" min="{{novel_upscale_ratio_min}}" max="{{novel_upscale_ratio_max}}" step="{{novel_upscale_ratio_step}}" value="{{novel_upscale_ratio}}" />
+            </div>
             <label for="sd_prompt_prefix">Common prompt prefix</label>
             <textarea id="sd_prompt_prefix" class="text_pole textarea_compact" rows="3"></textarea>
             <div id="sd_character_prompt_block">

public/scripts/nai-settings.js

@@ -1,7 +1,6 @@
 import {
     getRequestHeaders,
     getStoppingStrings,
-    max_context,
     novelai_setting_names,
     saveSettingsDebounced,
     setGenerationParamsFromPreset
@@ -15,13 +14,6 @@ import {
     uuidv4,
 } from "./utils.js";
 
-export {
-    nai_settings,
-    loadNovelPreset,
-    loadNovelSettings,
-    getNovelTier,
-};
-
 const default_preamble = "[ Style: chat, complex, sensory, visceral ]";
 const default_order = [1, 5, 0, 2, 3, 4];
 const maximum_output_length = 150;
@@ -32,7 +24,7 @@ const default_presets = {
     "kayra-v1": "Carefree-Kayra"
 }
 
-const nai_settings = {
+export const nai_settings = {
     temperature: 1.5,
     repetition_penalty: 2.25,
     repetition_penalty_range: 2048,
@@ -84,11 +76,33 @@ export function getKayraMaxContextTokens() {
     return null;
 }
 
-function getNovelTier(tier) {
+export function getNovelTier() {
-    return nai_tiers[tier] ?? 'no_connection';
+    return nai_tiers[novel_data?.tier] ?? 'no_connection';
 }
 
-function loadNovelPreset(preset) {
+export function getNovelAnlas() {
+    return novel_data?.trainingStepsLeft?.fixedTrainingStepsLeft ?? 0;
+}
+
+export function getNovelUnlimitedImageGeneration() {
+    return novel_data?.perks?.unlimitedImageGeneration ?? false;
+}
+
+export async function loadNovelSubscriptionData() {
+    const result = await fetch('/getstatus_novelai', {
+        method: 'POST',
+        headers: getRequestHeaders(),
+    });
+
+    if (result.ok) {
+        const data = await result.json();
+        setNovelData(data);
+    }
+
+    return result.ok;
+}
+
+export function loadNovelPreset(preset) {
     if (preset.genamt === undefined) {
         const needsUnlock = preset.max_context > MAX_CONTEXT_DEFAULT;
         $("#amount_gen").val(preset.max_length).trigger('input');
@@ -124,7 +138,7 @@ function loadNovelPreset(preset) {
     loadNovelSettingsUi(nai_settings);
 }
 
-function loadNovelSettings(settings) {
+export function loadNovelSettings(settings) {
     //load the rest of the Novel settings without any checks
     nai_settings.model_novel = settings.model_novel;
     $('#model_novel_select').val(nai_settings.model_novel);

server.js

@@ -4484,8 +4484,8 @@ app.post('/api/novelai/generate-image', jsonParser, async (request, response) =>
 
     try {
         console.log('NAI Diffusion request:', request.body);
-        const url = `${API_NOVELAI}/ai/generate-image`;
+        const generateUrl = `${API_NOVELAI}/ai/generate-image`;
-        const result = await fetch(url, {
+        const generateResult = await fetch(generateUrl, {
             method: 'POST',
             headers: {
                 'Authorization': `Bearer ${key}`,
@@ -4511,45 +4511,50 @@ app.post('/api/novelai/generate-image', jsonParser, async (request, response) =>
             }),
         });
 
-        if (!result.ok) {
+        if (!generateResult.ok) {
+            console.log('NovelAI returned an error.', generateResult.statusText);
             return response.sendStatus(500);
         }
 
-        const archiveBuffer = await result.arrayBuffer();
+        const archiveBuffer = await generateResult.arrayBuffer();
+        const imageBuffer = await extractFileFromZipBuffer(archiveBuffer, '.png');
+        const originalBase64 = imageBuffer.toString('base64');
 
-        const imageBuffer = await new Promise((resolve, reject) => yauzl.fromBuffer(Buffer.from(archiveBuffer), { lazyEntries: true }, (err, zipfile) => {
+        // No upscaling
-            if (err) {
+        if (isNaN(request.body.upscale_ratio) || request.body.upscale_ratio <= 1) {
-                reject(err);
+            return response.send(originalBase64);
         }
 
-            zipfile.readEntry();
+        try {
-            zipfile.on('entry', (entry) => {
+            console.debug('Upscaling image...');
-                if (entry.fileName.endsWith('.png')) {
+            const upscaleUrl = `${API_NOVELAI}/ai/upscale`;
-                    console.log(`Extracting ${entry.fileName}`);
+            const upscaleResult = await fetch(upscaleUrl, {
-                    zipfile.openReadStream(entry, (err, readStream) => {
+                method: 'POST',
-                        if (err) {
+                headers: {
-                            reject(err);
+                    'Authorization': `Bearer ${key}`,
-                        } else {
+                    'Content-Type': 'application/json',
-                            const chunks = [];
+                },
-                            readStream.on('data', (chunk) => {
+                body: JSON.stringify({
-                                chunks.push(chunk);
+                    image: originalBase64,
+                    height: request.body.height,
+                    width: request.body.width,
+                    scale: request.body.upscale_ratio,
+                }),
             });
 
-                            readStream.on('end', () => {
+            if (!upscaleResult.ok) {
-                                const buffer = Buffer.concat(chunks);
+                throw new Error('NovelAI returned an error.');
-                                resolve(buffer);
-                                zipfile.readEntry(); // Continue to the next entry
-                            });
             }
-                    });
-                } else {
-                    zipfile.readEntry(); // Continue to the next entry
-                }
-            });
-        }));
 
-        const base64 = imageBuffer.toString('base64');
+            const upscaledArchiveBuffer = await upscaleResult.arrayBuffer();
-        return response.send(base64);
+            const upscaledImageBuffer = await extractFileFromZipBuffer(upscaledArchiveBuffer, '.png');
+            const upscaledBase64 = upscaledImageBuffer.toString('base64');
+
+            return response.send(upscaledBase64);
+        } catch (error) {
+            console.warn('NovelAI generated an image, but upscaling failed. Returning original image.');
+            return response.send(originalBase64)
+        }
     } catch (error) {
         console.log(error);
         return response.sendStatus(500);
@@ -5103,6 +5108,45 @@ app.post('/import_custom', jsonParser, async (request, response) => {
     }
 });
 
+/**
+ * Extracts a file with given extension from an ArrayBuffer containing a ZIP archive.
+ * @param {ArrayBuffer} archiveBuffer Buffer containing a ZIP archive
+ * @param {string} fileExtension File extension to look for
+ * @returns {Promise<Buffer>} Buffer containing the extracted file
+ */
+async function extractFileFromZipBuffer(archiveBuffer, fileExtension) {
+    return await new Promise((resolve, reject) => yauzl.fromBuffer(Buffer.from(archiveBuffer), { lazyEntries: true }, (err, zipfile) => {
+        if (err) {
+            reject(err);
+        }
+
+        zipfile.readEntry();
+        zipfile.on('entry', (entry) => {
+            if (entry.fileName.endsWith(fileExtension)) {
+                console.log(`Extracting ${entry.fileName}`);
+                zipfile.openReadStream(entry, (err, readStream) => {
+                    if (err) {
+                        reject(err);
+                    } else {
+                        const chunks = [];
+                        readStream.on('data', (chunk) => {
+                            chunks.push(chunk);
+                        });
+
+                        readStream.on('end', () => {
+                            const buffer = Buffer.concat(chunks);
+                            resolve(buffer);
+                            zipfile.readEntry(); // Continue to the next entry
+                        });
+                    }
+                });
+            } else {
+                zipfile.readEntry();
+            }
+        });
+    }));
+}
+
 async function downloadChubLorebook(id) {
     const result = await fetch('https://api.chub.ai/api/lorebooks/download', {
         method: 'POST',