Move Stability generation to backend

e32b0cc223e3dc3d1895f51300aca4ff1660651e

Cohee <18619528+Cohee1207@users.noreply.github.com>

Parents: 6608e530c599632a61305a51788ec5d55cf16c9f
5 files changed, +198 -161Ignore whitespace
public/scripts/extensions/stable-diffusion/index.js+121 -109
@@ -22,7 +22,7 @@ import { getApiUrl, getContext, extension_settings, doExtrasFetch, modules, rend
2222import { selected_group } from '../../group-chats.js';
2323import { stringFormat, initScrollHeight, resetScrollHeight, getCharaFilename, saveBase64AsFile, getBase64Async, delay, isTrueBoolean, debounce } from '../../utils.js';
2424import { getMessageTimeStamp, humanizedDateTime } from '../../RossAscends-mods.js';
2525import { SECRET_KEYS, secret_state, writeSecret } from '../../secrets.js';
2626import { getNovelUnlimitedImageGeneration, getNovelAnlas, loadNovelSubscriptionData } from '../../nai-settings.js';
2727import { getMultimodalCaption } from '../shared.js';
2828import { SlashCommandParser } from '../../slash-commands/SlashCommandParser.js';
@@ -285,11 +285,7 @@ const defaultSettings = {
285285 interactive_visible: false,
286286
287287 // Stability AI settings
288288 stability_api_keystability_style_preset: 'anime',
289- stability_engine: 'V2beta Image Generation',
290- stability_style_preset: "anime",
291- stability_aspect_ratio: '1:1',
292- stability_output_format: 'png',
293289};
294290
295291const writePromptFieldsDebounced = debounce(writePromptFields, debounce_timeout.relaxed);
@@ -452,11 +448,7 @@ async function loadSettings() {
452448 $('#sd_wand_visible').prop('checked', extension_settings.sd.wand_visible);
453449 $('#sd_command_visible').prop('checked', extension_settings.sd.command_visible);
454450 $('#sd_interactive_visible').prop('checked', extension_settings.sd.interactive_visible);
455- $('#sd_stability_key').val(extension_settings.sd.stability_key);
456- $('#sd_stability_engine').val(extension_settings.sd.stability_engine);
457451 $('#sd_stability_style_preset').val(extension_settings.sd.stability_style_preset);
458- $('#sd_stability_aspect_ratio').val(extension_settings.sd.stability_aspect_ratio);
459- $('#sd_stability_output_format').val(extension_settings.sd.stability_output_format);
460452
461453 for (const style of extension_settings.sd.styles) {
462454 const option = document.createElement('option');
@@ -684,7 +676,7 @@ async function refinePrompt(prompt, allowExpand, isNegative = false) {
684676 const refinedPrompt = await callGenericPopup(text + 'Press "Cancel" to abort the image generation.', POPUP_TYPE.INPUT, prompt.trim(), { rows: 5, okButton: 'Continue' });
685677
686678 if (refinedPrompt) {
687679 return String(refinedPrompt);
688680 } else {
689681 throw new Error('Generation aborted by user.');
690682 }
@@ -1097,32 +1089,26 @@ function onComfyWorkflowChange() {
10971089 extension_settings.sd.comfy_workflow = $('#sd_comfy_workflow').find(':selected').val();
10981090 saveSettingsDebounced();
10991091}
1100-function onStabilityKeyInput() {
1101- extension_settings.sd.stability_key = $('#sd_stability_key').val();
1102- saveSettingsDebounced();
1103-}
11041092
11051093async function onStabilityEngineChangeonStabilityKeyClick() {
1106- extension_settings.sd.stability_engine = $('#sd_stability_engine').val();
1094+ const popupText = 'Stability AI API Key:';
1107- saveSettingsDebounced();
1095+ const key = await callGenericPopup(popupText, POPUP_TYPE.INPUT);
1108-}
11091096
11101097function if onStabilityStylePresetChange(!key) {
1111- extension_settings.sd.stability_style_preset = $('#sd_stability_style_preset').val();
1098+ return;
1112- saveSettingsDebounced();
1099+ }
1113-}
11141100
1115-function onStabilityAspectRatioChange() {
1101+ await writeSecret(SECRET_KEYS.STABILITY, String(key));
1116- extension_settings.sd.stability_aspect_ratio = $('#sd_stability_aspect_ratio').val();
1102+
1117- saveSettingsDebounced();
1103+ toastr.success('API Key saved');
1104+ await loadSettingOptions();
11181105}
11191106
11201107function onStabilityOutputFormatChangeonStabilityStylePresetChange() {
11211108 extension_settings.sd.stability_output_formatstability_style_preset = $('#sd_stability_output_formatsd_stability_style_preset').val();
11221109 saveSettingsDebounced();
11231110}
11241111
1125-
11261112async function changeComfyWorkflow(_, name) {
11271113 name = name.replace(/(\.json)?$/i, '.json');
11281114 if ($(`#sd_comfy_workflow > [value="${name}"]`).length > 0) {
@@ -1441,6 +1427,9 @@ async function loadSamplers() {
14411427 case sources.pollinations:
14421428 samplers = ['N/A'];
14431429 break;
1430+ case sources.stability:
1431+ samplers = ['N/A'];
1432+ break;
14441433 }
14451434
14461435 for (const sampler of samplers) {
@@ -1643,80 +1632,9 @@ async function loadModels() {
16431632 }
16441633}
16451634
1646-async function generateStabilityImage(prompt, negativePrompt) {
1647- const payload = {
1648- prompt: prompt,
1649- negative_prompt: negativePrompt,
1650- width: extension_settings.sd.width,
1651- height: extension_settings.sd.height,
1652- seed: extension_settings.sd.seed >= 0 ? extension_settings.sd.seed : undefined,
1653- style_preset: extension_settings.sd.stability_style_preset,
1654- output_format: extension_settings.sd.stability_output_format,
1655- };
1656-
1657- const formData = new FormData();
1658- for (const [key, value] of Object.entries(payload)) {
1659- if (value !== undefined) {
1660- formData.append(key, String(value));
1661- }
1662- }
1663-
1664- let apiUrl;
1665- switch (extension_settings.sd.model) {
1666- case 'stable-image-ultra':
1667- apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/ultra';
1668- break;
1669- case 'stable-image-core':
1670- apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/core';
1671- break;
1672- case 'stable-diffusion-3':
1673- apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/sd3';
1674- break;
1675- default:
1676- throw new Error('Invalid Stability AI model selected');
1677- }
1678-
1679- try {
1680- const response = await fetch(apiUrl, {
1681- method: 'POST',
1682- headers: {
1683- 'Authorization': `Bearer ${extension_settings.sd.stability_key}`,
1684- 'Accept': 'image/*',
1685- },
1686- body: formData,
1687- });
1688-
1689- if (!response.ok) {
1690- const errorText = await response.text();
1691- throw new Error(`HTTP ${response.status}: ${errorText}`);
1692- }
1693-
1694- const arrayBuffer = await response.arrayBuffer();
1695- const base64Image = arrayBufferToBase64(arrayBuffer);
1696-
1697- return {
1698- format: extension_settings.sd.stability_output_format,
1699- data: base64Image,
1700- };
1701- } catch (error) {
1702- console.error('Error generating image with Stability AI:', error);
1703- throw error;
1704- }
1705-}
1706-
1707-function arrayBufferToBase64(buffer) {
1708- let binary = '';
1709- const bytes = new Uint8Array(buffer);
1710- const len = bytes.byteLength;
1711- for (let i = 0; i < len; i++) {
1712- binary += String.fromCharCode(bytes[i]);
1713- }
1714- return btoa(binary);
1715-}
1716-
1717-
1718-
17191635async function loadStabilityModels() {
1636+ $('#sd_stability_key').toggleClass('success', !!secret_state[SECRET_KEYS.STABILITY]);
1637+
17201638 return [
17211639 { value: 'stable-image-ultra', text: 'Stable Image Ultra' },
17221640 { value: 'stable-image-core', text: 'Stable Image Core' },
@@ -2055,6 +1973,9 @@ async function loadSchedulers() {
20551973 case sources.comfy:
20561974 schedulers = await loadComfySchedulers();
20571975 break;
1976+ case sources.stability:
1977+ schedulers = ['N/A'];
1978+ break;
20581979 }
20591980
20601981 for (const scheduler of schedulers) {
@@ -2128,6 +2049,9 @@ async function loadVaes() {
21282049 case sources.comfy:
21292050 vaes = await loadComfyVaes();
21302051 break;
2052+ case sources.stability:
2053+ vaes = ['N/A'];
2054+ break;
21312055 }
21322056
21332057 for (const vae of vaes) {
@@ -2611,7 +2535,6 @@ async function sendGenerationRequest(generationType, prompt, additionalNegativeP
26112535 case sources.stability:
26122536 result = await generateStabilityImage(prefixedPrompt, negativePrompt);
26132537 break;
2614-
26152538 }
26162539
26172540 if (!result.data) {
@@ -2635,6 +2558,12 @@ async function sendGenerationRequest(generationType, prompt, additionalNegativeP
26352558 return base64Image;
26362559}
26372560
2561+/**
2562+ * Generates an image using the TogetherAI API.
2563+ * @param {string} prompt - The main instruction used to guide the image generation.
2564+ * @param {string} negativePrompt - The instruction used to restrict the image generation.
2565+ * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete.
2566+ */
26382567async function generateTogetherAIImage(prompt, negativePrompt) {
26392568 const result = await fetch('/api/sd/together/generate', {
26402569 method: 'POST',
@@ -2659,6 +2588,12 @@ async function generateTogetherAIImage(prompt, negativePrompt) {
26592588 }
26602589}
26612590
2591+/**
2592+ * Generates an image using the Pollinations API.
2593+ * @param {string} prompt - The main instruction used to guide the image generation.
2594+ * @param {string} negativePrompt - The instruction used to restrict the image generation.
2595+ * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete.
2596+ */
26622597async function generatePollinationsImage(prompt, negativePrompt) {
26632598 const result = await fetch('/api/sd/pollinations/generate', {
26642599 method: 'POST',
@@ -2728,6 +2663,86 @@ async function generateExtrasImage(prompt, negativePrompt) {
27282663}
27292664
27302665/**
2666+ * Gets an aspect ratio for Stability that is the closest to the given width and height.
2667+ * @param {number} width Target width
2668+ * @param {number} height Target height
2669+ * @returns {string} Closest aspect ratio as a string
2670+ */
2671+function getClosestAspectRatio(width, height) {
2672+ const aspectRatios = {
2673+ '16:9': 16 / 9,
2674+ '1:1': 1,
2675+ '21:9': 21 / 9,
2676+ '2:3': 2 / 3,
2677+ '3:2': 3 / 2,
2678+ '4:5': 4 / 5,
2679+ '5:4': 5 / 4,
2680+ '9:16': 9 / 16,
2681+ '9:21': 9 / 21,
2682+ };
2683+
2684+ const aspectRatio = width / height;
2685+
2686+ let closestAspectRatio = Object.keys(aspectRatios)[0];
2687+ let minDiff = Math.abs(aspectRatio - aspectRatios[closestAspectRatio]);
2688+
2689+ for (const key in aspectRatios) {
2690+ const diff = Math.abs(aspectRatio - aspectRatios[key]);
2691+ if (diff < minDiff) {
2692+ minDiff = diff;
2693+ closestAspectRatio = key;
2694+ }
2695+ }
2696+
2697+ return closestAspectRatio;
2698+}
2699+
2700+/**
2701+ * Generates an image using Stability AI.
2702+ * @param {string} prompt - The main instruction used to guide the image generation.
2703+ * @param {string} negativePrompt - The instruction used to restrict the image generation.
2704+ * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete.
2705+ */
2706+async function generateStabilityImage(prompt, negativePrompt) {
2707+ const IMAGE_FORMAT = 'png';
2708+ const PROMPT_LIMIT = 10000;
2709+
2710+ try {
2711+ const response = await fetch('/api/sd/stability/generate', {
2712+ method: 'POST',
2713+ headers: getRequestHeaders(),
2714+ body: JSON.stringify({
2715+ model: extension_settings.sd.model,
2716+ payload: {
2717+ prompt: prompt.slice(0, PROMPT_LIMIT),
2718+ negative_prompt: negativePrompt.slice(0, PROMPT_LIMIT),
2719+ aspect_ratio: getClosestAspectRatio(extension_settings.sd.width, extension_settings.sd.height),
2720+ seed: extension_settings.sd.seed >= 0 ? extension_settings.sd.seed : undefined,
2721+ style_preset: extension_settings.sd.stability_style_preset,
2722+ output_format: IMAGE_FORMAT,
2723+ },
2724+ }),
2725+ });
2726+
2727+ if (!response.ok) {
2728+ const errorText = await response.text();
2729+ throw new Error(`HTTP ${response.status}: ${errorText}`);
2730+ }
2731+
2732+ const blob = await response.blob();
2733+ const base64Image = await getBase64Async(blob);
2734+
2735+ return {
2736+ format: IMAGE_FORMAT,
2737+ data: base64Image,
2738+ };
2739+ } catch (error) {
2740+ console.error('Error generating image with Stability AI:', error);
2741+ throw error;
2742+ }
2743+}
2744+
2745+/**
27312746 * Generates a "horde" image using the provided prompt and configuration settings.
27322747 *
27332748 * @param {string} prompt - The main instruction used to guide the image generation.
@@ -3356,7 +3371,7 @@ function isValidState() {
33563371 case sources.pollinations:
33573372 return true;
33583373 case sources.stability:
3359- return !!extension_settings.sd.stability_key;
3374+ return secret_state[SECRET_KEYS.STABILITY];
33603375 }
33613376}
33623377
@@ -3584,12 +3599,9 @@ jQuery(async () => {
35843599 $('#sd_command_visible').on('input', onCommandVisibleInput);
35853600 $('#sd_interactive_visible').on('input', onInteractiveVisibleInput);
35863601 $('#sd_swap_dimensions').on('click', onSwapDimensionsClick);
35873602 $('#sd_stability_key').on('inputclick', onStabilityKeyInputonStabilityKeyClick);
3588- $('#sd_stability_engine').on('change', onStabilityEngineChange);
35893603 $('#sd_stability_style_preset').on('change', onStabilityStylePresetChange);
3590- $('#sd_stability_aspect_ratio').on('change', onStabilityAspectRatioChange);
3604+
3591- $('#sd_stability_output_format').on('change', onStabilityOutputFormatChange);
3592-
35933605 $('.sd_settings .inline-drawer-toggle').on('click', function () {
35943606 initScrollHeight($('#sd_prompt_prefix'));
35953607 initScrollHeight($('#sd_negative_prompt'));
public/scripts/extensions/stable-diffusion/settings.html+15 -52
@@ -44,10 +44,10 @@
4444 <option value="openai">OpenAI (DALL-E)</option>
4545 <option value="pollinations">Pollinations</option>
4646 <option value="vlad">SD.Next (vladmandic)</option>
47+ <option value="stability">Stability AI</option>
4748 <option value="auto">Stable Diffusion Web UI (AUTOMATIC1111)</option>
4849 <option value="horde">Stable Horde</option>
4950 <option value="togetherai">TogetherAI</option>
50- <option value="stability">Stability AI</option>
5151 </select>
5252 <div data-sd-source="auto">
5353 <label for="sd_auto_url">SD Web UI URL</label>
@@ -191,27 +191,22 @@
191191 </div>
192192 </div>
193193 <div data-sd-source="stability">
194- <label for="sd_stability_key">API Key</label>
194+ <div class="flex-container flexnowrap alignItemsBaseline marginBot5">
195- <div class="flex-container flexnowrap">
195+ <strong class="flex1" data-i18n="API Key">API Key</strong>
196196 <inputdiv id="sd_stability_key" type="password" class="text_pole flex1" placeholder="Enter your Stability AI APImenu_button keymenu_button_icon" />
197- <div id="sd_stability_validate" class="menu_button menu_button_icon">
197+ <i class="fa-fw fa-solid fa-key"></i>
198- <i class="fa-solid fa-check"></i>
198+ <span data-i18n="Click to set">Click to set</span>
199- <span data-i18n="Connect">
200- Connect
201- </span>
202199 </div>
203200 </div>
204- <i>You can find your API key in the Stability AI dashboard.</i>
201+ <div class="marginBot5">
205-
202+ <i data-i18n="You can find your API key in the Stability AI dashboard.">
203+ You can find your API key in the Stability AI dashboard.
204+ </i>
205+ </div>
206+
206207 <div class="flex-container">
207208 <div class="flex1">
208209 <label for="sd_stability_enginesd_stability_style_preset" data-i18n="Style Preset">EngineStyle Preset</label>
209- <select id="sd_stability_engine">
210- <option value="v2beta">V2beta Image Generation</option>
211- </select>
212- </div>
213- <div class="flex1">
214- <label for="sd_stability_style_preset">Style Preset</label>
215210 <select id="sd_stability_style_preset">
216211 <option value="anime">Anime</option>
217212 <option value="3d-model">3D Model</option>
@@ -233,39 +228,7 @@
233228 </select>
234229 </div>
235230 </div>
236-
231+ </div>
237- <div class="flex-container">
238- <div class="flex1">
239- <label for="sd_stability_aspect_ratio">Aspect Ratio</label>
240- <select id="sd_stability_aspect_ratio">
241- <option value="16:9">16:9</option>
242- <option value="1:1">1:1</option>
243- <option value="21:9">21:9</option>
244- <option value="2:3">2:3</option>
245- <option value="3:2">3:2</option>
246- <option value="4:5">4:5</option>
247- <option value="5:4">5:4</option>
248- <option value="9:16">9:16</option>
249- <option value="9:21">9:21</option>
250- </select>
251- </div>
252- <div class="flex1">
253- <label for="sd_stability_seed">Seed</label>
254- <input id="sd_stability_seed" type="number" class="text_pole" value="0" min="0" max="4294967295" />
255- </div>
256- </div>
257-
258- <div class="flex-container">
259- <div class="flex1">
260- <label for="sd_stability_output_format">Output Format</label>
261- <select id="sd_stability_output_format">
262- <option value="png">PNG</option>
263- <option value="webp">WebP</option>
264- <option value="jpeg">JPEG</option>
265- </select>
266- </div>
267- </div>
268- </div>
269232 <div class="flex-container">
270233 <div class="flex1">
271234 <label for="sd_model" data-i18n="Model">Model</label>
@@ -415,7 +378,7 @@
415378 </label>
416379 </div>
417380
418381 <div data-sd-source="novel,togetherai,pollinations,comfy,drawthings,vlad,auto,horde,extras,stability" class="marginTop5">
419382 <label for="sd_seed">
420383 <span data-i18n="Seed">Seed</span>
421384 <small data-i18n="(-1 for random)">(-1 for random)</small>
public/scripts/secrets.js+1 -0
@@ -31,6 +31,7 @@ export const SECRET_KEYS = {
3131 FEATHERLESS: 'api_key_featherless',
3232 ZEROONEAI: 'api_key_01ai',
3333 HUGGINGFACE: 'api_key_huggingface',
34+ STABILITY: 'api_key_stability',
3435};
3536
3637const INPUT_MAP = {
src/endpoints/secrets.js+1 -0
@@ -43,6 +43,7 @@ const SECRET_KEYS = {
4343 FEATHERLESS: 'api_key_featherless',
4444 ZEROONEAI: 'api_key_01ai',
4545 HUGGINGFACE: 'api_key_huggingface',
46+ STABILITY: 'api_key_stability',
4647};
4748
4849// These are the keys that are safe to expose, even if allowKeysExposure is false
src/endpoints/stable-diffusion.js+60 -0
@@ -7,6 +7,7 @@ const path = require('path');
77const writeFileAtomicSync = require('write-file-atomic').sync;
88const { jsonParser } = require('../express-common');
99const { readSecret, SECRET_KEYS } = require('./secrets.js');
10+const FormData = require('form-data');
1011
1112/**
1213 * Sanitizes a string.
@@ -793,9 +794,68 @@ pollinations.post('/generate', jsonParser, async (request, response) => {
793794 }
794795});
795796
797+const stability = express.Router();
798+
799+stability.post('/generate', jsonParser, async (request, response) => {
800+ try {
801+ const key = readSecret(request.user.directories, SECRET_KEYS.STABILITY);
802+
803+ if (!key) {
804+ console.log('Stability AI key not found.');
805+ return response.sendStatus(400);
806+ }
807+
808+ const { payload, model } = request.body;
809+
810+ const formData = new FormData();
811+ for (const [key, value] of Object.entries(payload)) {
812+ if (value !== undefined) {
813+ formData.append(key, String(value));
814+ }
815+ }
816+
817+ let apiUrl;
818+ switch (model) {
819+ case 'stable-image-ultra':
820+ apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/ultra';
821+ break;
822+ case 'stable-image-core':
823+ apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/core';
824+ break;
825+ case 'stable-diffusion-3':
826+ apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/sd3';
827+ break;
828+ default:
829+ throw new Error('Invalid Stability AI model selected');
830+ }
831+
832+ const result = await fetch(apiUrl, {
833+ method: 'POST',
834+ headers: {
835+ 'Authorization': `Bearer ${key}`,
836+ 'Accept': 'image/*',
837+ },
838+ body: formData,
839+ });
840+
841+ if (!result.ok) {
842+ const text = await result.text();
843+ console.log('Stability AI returned an error.', result.status, result.statusText, text);
844+ return response.sendStatus(500);
845+ }
846+
847+ const buffer = await result.buffer();
848+ return response.send(buffer);
849+ } catch (error) {
850+ console.log(error);
851+ return response.sendStatus(500);
852+ }
853+});
854+
796855router.use('/comfy', comfy);
797856router.use('/together', together);
798857router.use('/drawthings', drawthings);
799858router.use('/pollinations', pollinations);
859+router.use('/stability', stability);
800860
801861module.exports = { router };