From 623453d4348978ca816b7e700f879a07a177fe2a Mon Sep 17 00:00:00 2001
From: Banyanael <nbanyan@gmail.com>
Date: Sat, 11 Nov 2023 01:37:11 -0500
Subject: [PATCH] Improving AI generation queuing to allow rendering
to occur in the background without interrupting gameplay.
* Added a game options setting to allow use of different generation
steps for temporary images. Personally, I use 30 steps for temporary
images and 300 for persistent images.
* Temporary render requests are added to the beginning of the queue and
use PassageSwitchHandler to cancel and remove the request if the player
moves to another passage before rendering is completed.
* Added an option to automatically regenerate images at regular intervals.
* Added "float: inline-start" to report images to prevent them cascading
oddly when the report descriptions are short.
* Added regular cleaning of the imageDB to discard images that no longer
have attribution.
* Added a button for issuing the Interrupt command to SD and clearing
the generation queue.
---
css/endWeek/slavesReport.css | 1 +
js/003-data/gameVariableData.js | 2 +
src/art/artJS.js | 6 +-
src/art/genAI/imageDB.js | 24 ++++
src/art/genAI/stableDiffusion.js | 181 +++++++++++++++++++++----------
src/gui/options/options.js | 48 +++++---
src/js/main.js | 14 +++
7 files changed, 202 insertions(+), 74 deletions(-)
diff --git a/css/endWeek/slavesReport.css b/css/endWeek/slavesReport.css
index 041e1f7c6ef..c7f8dbc7a3b 100644
--- a/css/endWeek/slavesReport.css
+++ b/css/endWeek/slavesReport.css
@@ -1,4 +1,5 @@
div.slave-report {
margin-top: 1em;
margin-bottom: 1em;
+ float: inline-start;
}
diff --git a/js/003-data/gameVariableData.js b/js/003-data/gameVariableData.js
index 2b59fede063..6af7e14cbf9 100644
--- a/js/003-data/gameVariableData.js
+++ b/js/003-data/gameVariableData.js
@@ -174,6 +174,7 @@ App.Data.defaultGameStateVariables = {
// Stable Diffusion settings
aiApiUrl: "http://localhost:7860",
aiAutoGen: true,
+ aiAutoGenFrequency: 0,
aiCfgScale: 5,
aiCustomImagePrompts: 0,
aiCustomStyleNeg: "",
@@ -183,6 +184,7 @@ App.Data.defaultGameStateVariables = {
aiNationality: 2,
aiSamplingMethod: "DPM++ 2M SDE Karras",
aiSamplingSteps: 20,
+ aiSamplingStepsEvent: 20,
aiStyle: 1,
aiRestoreFaces: true,
aiUpscale: false,
diff --git a/src/art/artJS.js b/src/art/artJS.js
index c381189a6ba..dbdda0dc7cd 100644
--- a/src/art/artJS.js
+++ b/src/art/artJS.js
@@ -677,14 +677,12 @@ App.Art.aiArtElement = function(slave, imageSize) {
makeImageNavigationArrows(container);
function updateAndRefresh(index = null) {
- const imageGenerator = new App.Art.GenAI.StableDiffusionClient();
-
container.classList.add("refreshing");
- imageGenerator.updateSlave(slave, index).then(() => {
+ App.Art.GenAI.client.updateSlave(slave, index).then(() => {
refresh();
}).catch(error => {
- console.error(error);
+ console.log(error.message || error);
}).finally(() => {
container.classList.remove("refreshing");
});
diff --git a/src/art/genAI/imageDB.js b/src/art/genAI/imageDB.js
index 71868cf5f37..0775d76beb8 100644
--- a/src/art/genAI/imageDB.js
+++ b/src/art/genAI/imageDB.js
@@ -126,6 +126,29 @@ App.Art.GenAI.imageDB = (function() {
});
}
+ /**
+ * Purge all images from DB that don't have an associated slave
+ */
+ async function clean() {
+ await waitForInit();
+ return new Promise((resolve, reject) => {
+ let transaction = db.transaction(['images'], 'readwrite');
+ let objectStore = transaction.objectStore('images');
+
+ let slaveKeys = V.slaves.reduce((t, s) => {return [...t, ...s.custom.aiImageIds]}, []);
+
+ const allKeysRequest = objectStore.getAllKeys();
+ allKeysRequest.onsuccess = () => {
+ allKeysRequest.result.forEach(k => {
+ if (!slaveKeys.includes(k)){
+ objectStore.delete(k);
+ }
+ });
+ resolve();
+ };
+ });
+ }
+
/**
* Count the images currently in the DB
* @returns {Promise<number>}
@@ -149,6 +172,7 @@ App.Art.GenAI.imageDB = (function() {
getImage,
removeImage,
clear,
+ clean,
count
};
})();
diff --git a/src/art/genAI/stableDiffusion.js b/src/art/genAI/stableDiffusion.js
index 41b2ec4f183..94925716f75 100644
--- a/src/art/genAI/stableDiffusion.js
+++ b/src/art/genAI/stableDiffusion.js
@@ -84,43 +84,58 @@ async function fetchWithTimeout(url, timeout, options) {
App.Art.GenAI.StableDiffusionClientQueue = class {
constructor() {
- /** @type {Array<{slaveID: number, body: string, resolve: function(Object): void, reject: function(string): void}>} */
+ /** @type {Array<{slaveID: number, body: string, resolves: [function(Object): void], rejects: [function(string): void]}>} */
this.queue = [];
this.interrupted = false;
+ this.workingOnID = false;
}
/**
* Process the top item in the queue, and continue processing the queue one at a time afterwards
* @private
*/
- async process() {
- while (this.queue.length > 0 && !this.interrupted) {
- const top = this.queue.first();
-
- // find all the requests for this slave in the queue; we'll satisfy them all at once
- // using only the newest data (i.e. the data from the last member)
- const satisfied = this.queue.filter(item => item.slaveID === top.slaveID);
- console.log(`Fetching image for slave ${top.slaveID}, satisfying ${satisfied.length} requests`);
-
+ process() {
+ if (this.workingOnID !== false) {
+ return false;
+ }
+ if (this.interrupted) {
+ return false;
+ }
+ const top = this.queue.shift();
+ if (!top) {
+ return false;
+ }
+ try {
+ this.workingOnID = top.slaveID;
+ console.log(`Fetching image for slave ${top.slaveID}, ${this.queue.length} requests remaining in the queue.`);
+ console.log("Generation Settings: ", JSON.parse(top.body));
const options = {
method: "POST",
headers: {
"Content-Type": "application/json",
},
- body: satisfied.last().body,
+ body: top.body,
};
- try {
- const response = await fetchWithTimeout(`${V.aiApiUrl}/sdapi/v1/txt2img`, 600000, options);
- if (!response.ok) {
- throw new Error(`Error fetching Stable Diffusion image - status: ${response.status}`);
- }
- const obj = await response.json();
- satisfied.forEach(item => item.resolve(obj));
- } catch (e) {
- satisfied.forEach(item => item.reject(e));
- }
- this.queue.delete(...satisfied);
+ const response = fetchWithTimeout(`${V.aiApiUrl}/sdapi/v1/txt2img`, 3000 * V.aiSamplingSteps, options)
+ .then((value) => {
+ value.json()
+ .then (obj => {
+ top.resolves.forEach(resolve => resolve(obj));
+ this.workingOnID = false;
+ this.process();
+ });
+ })
+ .catch(err => {
+ this.workingOnID = false;
+ top.rejects.forEach(reject => reject(`${top.slaveID}: Error fetching Stable Diffusion image - status: ${err}`));
+ this.process();
+ });
+ } catch (err) {
+ this.workingOnID = false;
+ top.rejects.forEach(reject => reject(err));
+ this.process();
}
+ return true;
}
/**
@@ -133,6 +148,16 @@ App.Art.GenAI.StableDiffusionClientQueue = class {
}
}
+ /**
+ * await this in order to block until the queue stops processing
+ */
+ async resumeAfterProcessing() {
+ const sleep = () => new Promise(r => setTimeout(r, 10));
+ while (this.workingOnID !== false) {
+ await sleep();
+ }
+ }
+
/**
* Queue image generation for an entity
* @param {number} slaveID or a unique negative value for non-slave entities
@@ -144,16 +169,36 @@ App.Art.GenAI.StableDiffusionClientQueue = class {
await this.resumeAfterInterrupt();
}
+ if (slaveID !== null) {
+ let item = this.queue.find(i => i.slaveID == slaveID);
+ if (item) {
+ // if id is already queued, add a handle to receive the previously queued Promise's response and update `body` with the new query
+ return new Promise((resolve, reject) => {
+ item.body = body;
+ item.resolves.push(resolve);
+ item.rejects.push(reject);
+ });
+ }
+ }
return new Promise((resolve, reject) => {
- this.queue.push({
- slaveID,
- body,
- resolve,
- reject
- });
- if (this.queue.length === 1) {
- this.process(); // do not await
+ if (App.Art.GenAI.client.renderEventImage()) {
+ // inject event images to the beginning of the queue
+ this.queue.unshift({
+ slaveID: slaveID,
+ body: body,
+ resolves: [resolve],
+ rejects: [reject]
+ });
+ } else {
+ this.queue.push({
+ slaveID: slaveID,
+ body: body,
+ resolves: [resolve],
+ rejects: [reject]
+ });
}
+
+ this.process(); // do not await
});
}
@@ -162,7 +207,7 @@ App.Art.GenAI.StableDiffusionClientQueue = class {
*/
async interrupt() {
if (this.interrupted) { // permit nesting and consecutive calls
- return;
+ return false;
}
this.interrupted = true; // pause processing of the queue and don't accept further interrupts
@@ -174,21 +219,23 @@ App.Art.GenAI.StableDiffusionClientQueue = class {
"Content-Type": "application/json",
},
};
- try {
- await fetchWithTimeout(`${V.aiApiUrl}/sdapi/v1/interrupt`, 1000, options);
- } catch {
- // ignore errors
- }
-
// reject everything in the queue
- for (const item of this.queue) {
- item.reject("Stable Diffusion fetch interrupted");
+ while (this.queue.length > 0) {
+ let item = this.queue.pop();
+ if (item) {
+ item.rejects.forEach(r => r(`${item.slaveID}: Stable Diffusion fetch interrupted`));
+ }
}
this.queue = [];
+
+ fetchWithTimeout(`${V.aiApiUrl}/sdapi/v1/interrupt`, 1000, options).then(() => {console.log("Stable Diffusion: Interrupt Sent.")}).catch (() => {
+ // ignore errors
+ });
this.interrupted = false; // resume with next add
+ return true;
}
-};
+}
// instantiate global queue
App.Art.GenAI.sdQueue = new App.Art.GenAI.StableDiffusionClientQueue();
@@ -210,11 +257,10 @@ App.Art.GenAI.StableDiffusionClient = class {
prompt: prompt.positive(),
sampler_name: V.aiSamplingMethod,
seed: slave.natural.artSeed,
- steps: V.aiSamplingSteps,
+ steps: this.renderEventImage() ? V.aiSamplingStepsEvent : V.aiSamplingSteps,
width: V.aiWidth,
restore_faces: V.aiRestoreFaces
});
-
return settings;
}
@@ -224,16 +270,26 @@ App.Art.GenAI.StableDiffusionClient = class {
*/
async fetchImageForSlave(slave) {
const settings = this.buildStableDiffusionSettings(slave);
- // set up a passage switch handler to interrupt image generation if it's incomplete
- const oldHandler = App.Utils.PassageSwitchHandler.get();
- App.Utils.PassageSwitchHandler.set(() => {
- App.Art.GenAI.sdQueue.interrupt();
- if (oldHandler) {
- oldHandler();
- }
- });
-
- const response = await App.Art.GenAI.sdQueue.add(slave.ID, JSON.stringify(settings));
+ const body = JSON.stringify(settings);
+ // set up a passage switch handler to clear queued generation of this event image upon passage change
+ if (this.renderEventImage()) {
+ const oldHandler = App.Utils.PassageSwitchHandler.get();
+ App.Utils.PassageSwitchHandler.set(() => {
+ // find where this request is in the queue
+ const rIndex = App.Art.GenAI.sdQueue.queue.findIndex(r => r.slaveID == slave.ID && r.body == body);
+ if (rIndex > -1) {
+ const rejects = App.Art.GenAI.sdQueue.queue[rIndex].rejects;
+ // remove request from the queue as soon as possible
+ App.Art.GenAI.sdQueue.queue.splice(rIndex, 1);
+ // reject the associated promises
+ rejects.forEach(r => r(`${slave.ID} (Event): Stable Diffusion fetch interrupted`));
+ }
+ if (oldHandler) {
+ oldHandler();
+ }
+ });
+ }
+ const response = await App.Art.GenAI.sdQueue.add(slave.ID, body);
return response.images[0];
}
@@ -243,28 +299,43 @@ App.Art.GenAI.StableDiffusionClient = class {
* @param {number | null} replacementImageIndex - If provided, replace the image at this index
*/
async updateSlave(slave, replacementImageIndex = null) {
- const base64Image = await this.fetchImageForSlave(slave);
+ const base64Image = await this.fetchImageForSlave(slave)
const imageData = getImageData(base64Image);
const imagePreexisting = await compareExistingImages(slave, imageData);
-
+ let vSlave = V.slaves.find(s=>s.ID == slave.ID);
// If new image, add or replace it in
if (imagePreexisting === -1) {
const imageId = await App.Art.GenAI.imageDB.putImage({data: imageData});
if (replacementImageIndex !== null) {
await App.Art.GenAI.imageDB.removeImage(slave.custom.aiImageIds[replacementImageIndex]);
slave.custom.aiImageIds[replacementImageIndex] = imageId;
+ if (vSlave) vSlave.custom.aiImageIds[replacementImageIndex] = imageId;
} else {
slave.custom.aiImageIds.push(imageId);
+ if (vSlave) vSlave.custom.aiImageIds.push(imageId);
slave.custom.aiDisplayImageIdx = slave.custom.aiImageIds.indexOf(imageId);
+ if (vSlave) vSlave.custom.aiDisplayImageIdx = slave.custom.aiImageIds.indexOf(imageId);
}
// If image already exists, just update the display idx to it
} else {
- console.log('generated redundant image, no image stored');
+ console.log('Generated redundant image, no image stored');
slave.custom.aiDisplayImageIdx = imagePreexisting;
+ if (vSlave) vSlave.custom.aiDisplayImageIdx = imagePreexisting;
}
}
+
+ renderEventImage() {
+ // data-tags attribute values that identify pages where temporary images generation will be requested
+ const tags = ["end-week"];
+ // data-passage attribute values that identify pages where temporary images generation will be requested
+ const passages = ["Summary Options", "Buy Slaves", "Market", "Bulk Slave Intro", "Wardrobe"];
+ return tags.some(i => $(`[data-tags='${i}']`).length > 0) ||
+ passages.some(i => $(`[data-passage='${i}']`).length > 0);
+ }
};
+App.Art.GenAI.client = new App.Art.GenAI.StableDiffusionClient();
+
/**
* Search slave's existing images for a match with the new image.
* @param {FC.SlaveState} slave - The slave we're updating
diff --git a/src/gui/options/options.js b/src/gui/options/options.js
index e9b447e369d..40302364402 100644
--- a/src/gui/options/options.js
+++ b/src/gui/options/options.js
@@ -1277,12 +1277,18 @@ App.UI.artOptions = function() {
options.addOption("Automatic generation", "aiAutoGen")
.addValue("Enabled", true).on().addValue("Disabled", false).off()
.addComment("Generate images for new slaves on the fly. If disabled, you will need to manually click to generate each slave's image.");
+ if (V.aiAutoGen) {
+ options.addOption("Regeneration Frequency", "aiAutoGenFrequency").showTextBox()
+ .addComment("How often (in weeks) regenerate slave images. Set to 0 to disable. Slaves will render when 'Weeks Owned' is divisible by this number.");
+ }
options.addOption("Sampling Method", "aiSamplingMethod").showTextBox()
.addComment(`The sampling method used by AI. You can query ${V.aiApiUrl}/sdapi/v1/samplers to see the list of available samplers.`);
options.addOption("CFG Scale", "aiCfgScale").showTextBox()
.addComment("The higher this number, the more the prompt influences the image. Generally between 5 to 12.");
options.addOption("Sampling Steps", "aiSamplingSteps").showTextBox()
- .addComment("The number of steps used when generating the image. More steps might reduce artifacts but increases generation time. Generally between 20 to 50.");
+ .addComment("The number of steps used when generating the image. More steps might reduce artifacts but increases generation time. Generally between 20 to 50, but may be as high as 500 if you don't mind long queues in the background.");
+ options.addOption("Event Sampling Steps", "aiSamplingStepsEvent").showTextBox()
+ .addComment("The number of steps used when generating an image durring events. Generally between 20 to 50 to maintain a reasonable speed.");
options.addOption("Height", "aiHeight").showTextBox()
.addComment("The height of the image.");
options.addOption("Width", "aiWidth").showTextBox()
@@ -1299,25 +1305,37 @@ App.UI.artOptions = function() {
options.addOption("Upscaling method", "aiUpscaler").showTextBox()
.addComment(`The method used for upscaling the image. You can query ${V.aiApiUrl}/sdapi/v1/upscalers to see the list of available upscalers.`);
}
+ async function renderQueueOption(clicked = false){
+ const sleep = (ms) => new Promise(r => setTimeout(r, ms));
+ // wait for the button to render
+ while (!$("button:contains('Interrupt rendering')").length) {
+ await sleep(10);
+ }
+ if (clicked) {
+ // send interrupt when clicked
+ App.Art.GenAI.sdQueue.interrupt();
+ }
+ if (App.Art.GenAI.sdQueue.interrupted) {
+ $("button:contains('Interrupt rendering')").removeClass("off").addClass("on selected disabled");
+ await App.Art.GenAI.sdQueue.resumeAfterInterrupt();
+ }
+ $("button:contains('Interrupt rendering')").removeClass("on selected disabled").addClass("off");
+ }
+ options.addCustomOption("Rendering Queue management")
+ .addButton("Interrupt rendering", () => renderQueueOption(true));
+ // adjust the state of the button when it is rendered
+ renderQueueOption();
options.addCustomOption("Cache database management")
.addButton("Purge all images", async () => {
await App.Art.GenAI.imageDB.clear();
})
.addButton("Regenerate images for all slaves", () => {
- const lockID = LoadScreen.lock();
- const waitMessage = $(`<div class="endweek-titleblock"><div class="endweek-maintitle">Regenerating images for ${V.slaves.length} slaves...</div></div>`);
- $("#init-screen").append(waitMessage);
- // queue actual regeneration to happen *after* the passage change
- const regenFunc = () => {
- const generator = new App.Art.GenAI.StableDiffusionClient();
- Promise
- .all(V.slaves.map(s => generator.updateSlave(s)))
- .finally(() => {
- LoadScreen.unlock(lockID);
- waitMessage.remove();
- });
- };
- setTimeout(regenFunc, 50);
+ // queue all slaves for regeneration in the background
+ V.slaves.forEach(s => App.Art.GenAI.client.updateSlave(s)
+ .catch(error => {
+ console.log(error.message || error);
+ }));
+ console.log(`${App.Art.GenAI.sdQueue.queue.length} requests queued for rendering.`);
})
.addComment(`Current cache size: <span id="cacheCount">Please wait...</span> images. The cache database is shared between games.`);
App.Art.GenAI.imageDB.count().then((result) => {
diff --git a/src/js/main.js b/src/js/main.js
index 113780c21ff..4e20c7787d3 100644
--- a/src/js/main.js
+++ b/src/js/main.js
@@ -191,6 +191,20 @@ App.MainView.full = function() {
V.tabChoice.SlaveInteract = "Description";
}
+ // remove non-slave images from the DB to free up storage
+ App.Art.GenAI.imageDB.clean();
+
+ // regenerate old slave images
+ if (V.aiAutoGenFrequency > 0){
+ V.slaves.forEach(s=>{
+ if ((V.week - s.weekAcquired) % V.aiAutoGenFrequency == 0){
+ App.Art.GenAI.client.updateSlave(s)
+ .catch(error => {
+ console.log(error.message || error);
+ });
+ }
+ });
+ }
penthouseCensus();
V.costs = Math.trunc(calculateCosts.predict());
if (V.defaultRules.length > 0) {
--
GitLab