rixty/SillyTavern
1
0
Fork 0
mirror of https://github.com/SillyTavern/SillyTavern.git synced 2025-02-09 08:38:53 +01:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Cohee 2023-09-04 18:00:15 +03:00
parent e616ab5ced
commit ded1e3a859
5 changed files with 236 additions and 79 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
public/script.js
View File

@ -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({
type: "POST", //
url: "/getstatus_novelai", //
data: JSON.stringify(data),
beforeSend: function () {
if (!result) {
throw new Error('Could not load subscription data');
}
},
cache: false,
dataType: "json",
contentType: "application/json",
success: function (data) {
if (data.error != true) {
setNovelData(data);
online_status = `${getNovelTier(data.tier)}`;
}
resultCheckStatusNovel();
},
error: function (jqXHR, exception) {
online_status = "no_connection";
console.log(exception);
console.log(jqXHR);
resultCheckStatusNovel();
},
});
online_status = getNovelTier();
} catch {
online_status = "no_connection";
}
resultCheckStatusNovel();
} else {
if (is_get_status != true && is_get_status_openai != true) {
online_status = "no_connection";

108
public/scripts/extensions/stable-diffusion/index.js
View File

@ -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,
height: extension_settings.sd.height,
width: width,
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 () {

16
public/scripts/extensions/stable-diffusion/settings.html
View File

@ -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">

40
public/scripts/nai-settings.js
View File

@ -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) {
return nai_tiers[tier] ?? 'no_connection';
export function getNovelTier() {
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);

112
server.js
View File

@ -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 result = await fetch(url, {
const generateUrl = `${API_NOVELAI}/ai/generate-image`;
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) => {
if (err) {
reject(err);
// No upscaling
if (isNaN(request.body.upscale_ratio) || request.body.upscale_ratio <= 1) {
return response.send(originalBase64);
}
try {
console.debug('Upscaling image...');
const upscaleUrl = `${API_NOVELAI}/ai/upscale`;
const upscaleResult = await fetch(upscaleUrl, {
method: 'POST',
headers: {
'Authorization': `Bearer ${key}`,
'Content-Type': 'application/json',
},
body: JSON.stringify({
image: originalBase64,
height: request.body.height,
width: request.body.width,
scale: request.body.upscale_ratio,
}),
});
if (!upscaleResult.ok) {
throw new Error('NovelAI returned an error.');
}
zipfile.readEntry();
zipfile.on('entry', (entry) => {
if (entry.fileName.endsWith('.png')) {
console.log(`Extracting ${entry.fileName}`);
zipfile.openReadStream(entry, (err, readStream) => {
if (err) {
reject(err);
} else {
const chunks = [];
readStream.on('data', (chunk) => {
chunks.push(chunk);
});
const upscaledArchiveBuffer = await upscaleResult.arrayBuffer();
const upscaledImageBuffer = await extractFileFromZipBuffer(upscaledArchiveBuffer, '.png');
const upscaledBase64 = upscaledImageBuffer.toString('base64');
readStream.on('end', () => {
const buffer = Buffer.concat(chunks);
resolve(buffer);
zipfile.readEntry(); // Continue to the next entry
});
}
});
} else {
zipfile.readEntry(); // Continue to the next entry
}
});
}));
const base64 = imageBuffer.toString('base64');
return response.send(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',