Move Stability generation to backend

e32b0cc223e3dc3d1895f51300aca4ff1660651e

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

Parents: 6608e530c599632a61305a51788ec5d55cf16c9f

public/scripts/extensions/stable-diffusion/index.js+121 -109 Viewed

 import { selected_group } from '../../group-chats.js';
 import { stringFormat, initScrollHeight, resetScrollHeight, getCharaFilename, saveBase64AsFile, getBase64Async, delay, isTrueBoolean, debounce } from '../../utils.js';
 import { getMessageTimeStamp, humanizedDateTime } from '../../RossAscends-mods.js';
-import { SECRET_KEYS, secret_state } from '../../secrets.js';
+import { SECRET_KEYS, secret_state, writeSecret } from '../../secrets.js';
 import { getNovelUnlimitedImageGeneration, getNovelAnlas, loadNovelSubscriptionData } from '../../nai-settings.js';
 import { getMultimodalCaption } from '../shared.js';
 import { SlashCommandParser } from '../../slash-commands/SlashCommandParser.js';

     interactive_visible: false,
     // Stability AI settings
-    stability_api_key: '',
+    stability_style_preset: 'anime',
-    stability_engine: 'V2beta Image Generation',
-    stability_style_preset: "anime",
-    stability_aspect_ratio: '1:1',
-    stability_output_format: 'png',
 };
 const writePromptFieldsDebounced = debounce(writePromptFields, debounce_timeout.relaxed);

     $('#sd_wand_visible').prop('checked', extension_settings.sd.wand_visible);
     $('#sd_command_visible').prop('checked', extension_settings.sd.command_visible);
     $('#sd_interactive_visible').prop('checked', extension_settings.sd.interactive_visible);
-    $('#sd_stability_key').val(extension_settings.sd.stability_key);
-    $('#sd_stability_engine').val(extension_settings.sd.stability_engine);
     $('#sd_stability_style_preset').val(extension_settings.sd.stability_style_preset);
-    $('#sd_stability_aspect_ratio').val(extension_settings.sd.stability_aspect_ratio);
-    $('#sd_stability_output_format').val(extension_settings.sd.stability_output_format);
     for (const style of extension_settings.sd.styles) {
         const option = document.createElement('option');

         const refinedPrompt = await callGenericPopup(text + 'Press "Cancel" to abort the image generation.', POPUP_TYPE.INPUT, prompt.trim(), { rows: 5, okButton: 'Continue' });
         if (refinedPrompt) {
-            return refinedPrompt;
+            return String(refinedPrompt);
         } else {
             throw new Error('Generation aborted by user.');
+        }

     extension_settings.sd.comfy_workflow = $('#sd_comfy_workflow').find(':selected').val();
     saveSettingsDebounced();
+}
-function onStabilityKeyInput() {
-    extension_settings.sd.stability_key = $('#sd_stability_key').val();
-    saveSettingsDebounced();
-}
-function onStabilityEngineChange() {
+async function onStabilityKeyClick() {
-    extension_settings.sd.stability_engine = $('#sd_stability_engine').val();
+    const popupText = 'Stability AI API Key:';
-    saveSettingsDebounced();
+    const key = await callGenericPopup(popupText, POPUP_TYPE.INPUT);
-}
-function onStabilityStylePresetChange() {
+    if (!key) {
-    extension_settings.sd.stability_style_preset = $('#sd_stability_style_preset').val();
+        return;
-    saveSettingsDebounced();
+    }
-}
-function onStabilityAspectRatioChange() {
+    await writeSecret(SECRET_KEYS.STABILITY, String(key));
-    extension_settings.sd.stability_aspect_ratio = $('#sd_stability_aspect_ratio').val();
-    saveSettingsDebounced();
+    toastr.success('API Key saved');
+    await loadSettingOptions();
+}
-function onStabilityOutputFormatChange() {
+function onStabilityStylePresetChange() {
-    extension_settings.sd.stability_output_format = $('#sd_stability_output_format').val();
+    extension_settings.sd.stability_style_preset = $('#sd_stability_style_preset').val();
     saveSettingsDebounced();
+}
 async function changeComfyWorkflow(_, name) {
     name = name.replace(/(\.json)?$/i, '.json');
     if ($(`#sd_comfy_workflow > [value="${name}"]`).length > 0) {

         case sources.pollinations:
             samplers = ['N/A'];
             break;
+        case sources.stability:
+            samplers = ['N/A'];
+            break;
+    }
     for (const sampler of samplers) {

+    }
+}
-async function generateStabilityImage(prompt, negativePrompt) {
-    const payload = {
-        prompt: prompt,
-        negative_prompt: negativePrompt,
-        width: extension_settings.sd.width,
-        height: extension_settings.sd.height,
-        seed: extension_settings.sd.seed >= 0 ? extension_settings.sd.seed : undefined,
-        style_preset: extension_settings.sd.stability_style_preset,
-        output_format: extension_settings.sd.stability_output_format,
-    };
-    const formData = new FormData();
-    for (const [key, value] of Object.entries(payload)) {
-        if (value !== undefined) {
-            formData.append(key, String(value));
-        }
-    }
-    let apiUrl;
-    switch (extension_settings.sd.model) {
-        case 'stable-image-ultra':
-            apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/ultra';
-            break;
-        case 'stable-image-core':
-            apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/core';
-            break;
-        case 'stable-diffusion-3':
-            apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/sd3';
-            break;
-        default:
-            throw new Error('Invalid Stability AI model selected');
-    }
-    try {
-        const response = await fetch(apiUrl, {
-            method: 'POST',
-            headers: {
-                'Authorization': `Bearer ${extension_settings.sd.stability_key}`,
-                'Accept': 'image/*',
-            },
-            body: formData,
-        });
-        if (!response.ok) {
-            const errorText = await response.text();
-            throw new Error(`HTTP ${response.status}: ${errorText}`);
-        }
-        const arrayBuffer = await response.arrayBuffer();
-        const base64Image = arrayBufferToBase64(arrayBuffer);
-        return {
-            format: extension_settings.sd.stability_output_format,
-            data: base64Image,
-        };
-    } catch (error) {
-        console.error('Error generating image with Stability AI:', error);
-        throw error;
-    }
-}
-function arrayBufferToBase64(buffer) {
-    let binary = '';
-    const bytes = new Uint8Array(buffer);
-    const len = bytes.byteLength;
-    for (let i = 0; i < len; i++) {
-        binary += String.fromCharCode(bytes[i]);
-    }
-    return btoa(binary);
-}
 async function loadStabilityModels() {
+    $('#sd_stability_key').toggleClass('success', !!secret_state[SECRET_KEYS.STABILITY]);
     return [
         { value: 'stable-image-ultra', text: 'Stable Image Ultra' },
         { value: 'stable-image-core', text: 'Stable Image Core' },

         case sources.comfy:
             schedulers = await loadComfySchedulers();
             break;
+        case sources.stability:
+            schedulers = ['N/A'];
+            break;
+    }
     for (const scheduler of schedulers) {

         case sources.comfy:
             vaes = await loadComfyVaes();
             break;
+        case sources.stability:
+            vaes = ['N/A'];
+            break;
+    }
     for (const vae of vaes) {

             case sources.stability:
                 result = await generateStabilityImage(prefixedPrompt, negativePrompt);
                 break;
+        }
         if (!result.data) {

     return base64Image;
+}
+/**
+ * Generates an image using the TogetherAI API.
+ * @param {string} prompt - The main instruction used to guide the image generation.
+ * @param {string} negativePrompt - The instruction used to restrict the image generation.
+ * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete.
+ */
 async function generateTogetherAIImage(prompt, negativePrompt) {
     const result = await fetch('/api/sd/together/generate', {
         method: 'POST',

+    }
+}
+/**
+ * Generates an image using the Pollinations API.
+ * @param {string} prompt - The main instruction used to guide the image generation.
+ * @param {string} negativePrompt - The instruction used to restrict the image generation.
+ * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete.
+ */
 async function generatePollinationsImage(prompt, negativePrompt) {
     const result = await fetch('/api/sd/pollinations/generate', {
         method: 'POST',

+}
 /**
+ * Gets an aspect ratio for Stability that is the closest to the given width and height.
+ * @param {number} width Target width
+ * @param {number} height Target height
+ * @returns {string} Closest aspect ratio as a string
+ */
+function getClosestAspectRatio(width, height) {
+    const aspectRatios = {
+        '16:9': 16 / 9,
+        '1:1': 1,
+        '21:9': 21 / 9,
+        '2:3': 2 / 3,
+        '3:2': 3 / 2,
+        '4:5': 4 / 5,
+        '5:4': 5 / 4,
+        '9:16': 9 / 16,
+        '9:21': 9 / 21,
+    };
+    const aspectRatio = width / height;
+    let closestAspectRatio = Object.keys(aspectRatios)[0];
+    let minDiff = Math.abs(aspectRatio - aspectRatios[closestAspectRatio]);
+    for (const key in aspectRatios) {
+        const diff = Math.abs(aspectRatio - aspectRatios[key]);
+        if (diff < minDiff) {
+            minDiff = diff;
+            closestAspectRatio = key;
+        }
+    }
+    return closestAspectRatio;
+}
+/**
+ * Generates an image using Stability AI.
+ * @param {string} prompt - The main instruction used to guide the image generation.
+ * @param {string} negativePrompt - The instruction used to restrict the image generation.
+ * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete.
+ */
+async function generateStabilityImage(prompt, negativePrompt) {
+    const IMAGE_FORMAT = 'png';
+    const PROMPT_LIMIT = 10000;
+    try {
+        const response = await fetch('/api/sd/stability/generate', {
+            method: 'POST',
+            headers: getRequestHeaders(),
+            body: JSON.stringify({
+                model: extension_settings.sd.model,
+                payload: {
+                    prompt: prompt.slice(0, PROMPT_LIMIT),
+                    negative_prompt: negativePrompt.slice(0, PROMPT_LIMIT),
+                    aspect_ratio: getClosestAspectRatio(extension_settings.sd.width, extension_settings.sd.height),
+                    seed: extension_settings.sd.seed >= 0 ? extension_settings.sd.seed : undefined,
+                    style_preset: extension_settings.sd.stability_style_preset,
+                    output_format: IMAGE_FORMAT,
+                },
+            }),
+        });
+        if (!response.ok) {
+            const errorText = await response.text();
+            throw new Error(`HTTP ${response.status}: ${errorText}`);
+        }
+        const blob = await response.blob();
+        const base64Image = await getBase64Async(blob);
+        return {
+            format: IMAGE_FORMAT,
+            data: base64Image,
+        };
+    } catch (error) {
+        console.error('Error generating image with Stability AI:', error);
+        throw error;
+    }
+}
+/**
  * Generates a "horde" image using the provided prompt and configuration settings.
+ *
  * @param {string} prompt - The main instruction used to guide the image generation.

@@ -3356,7 +3371,7 @@ function isValidState() {
3356	case sources.pollinations:	3371	case sources.pollinations:
3357	return true;	3372	return true;
3358	case sources.stability:	3373	case sources.stability:
3359	return !!extension_settings.sd.stability_key;	3374	return secret_state[SECRET_KEYS.STABILITY];
3360	}	3375	}
3361	}	3376	}
3362		3377

     $('#sd_command_visible').on('input', onCommandVisibleInput);
     $('#sd_interactive_visible').on('input', onInteractiveVisibleInput);
     $('#sd_swap_dimensions').on('click', onSwapDimensionsClick);
-    $('#sd_stability_key').on('input', onStabilityKeyInput);
+    $('#sd_stability_key').on('click', onStabilityKeyClick);
-    $('#sd_stability_engine').on('change', onStabilityEngineChange);
     $('#sd_stability_style_preset').on('change', onStabilityStylePresetChange);
-    $('#sd_stability_aspect_ratio').on('change', onStabilityAspectRatioChange);
-    $('#sd_stability_output_format').on('change', onStabilityOutputFormatChange);
     $('.sd_settings .inline-drawer-toggle').on('click', function () {
         initScrollHeight($('#sd_prompt_prefix'));
         initScrollHeight($('#sd_negative_prompt'));

public/scripts/extensions/stable-diffusion/settings.html+15 -52 Viewed

                 <option value="openai">OpenAI (DALL-E)</option>
                 <option value="pollinations">Pollinations</option>
                 <option value="vlad">SD.Next (vladmandic)</option>
+                <option value="stability">Stability AI</option>
                 <option value="auto">Stable Diffusion Web UI (AUTOMATIC1111)</option>
                 <option value="horde">Stable Horde</option>
                 <option value="togetherai">TogetherAI</option>
-                <option value="stability">Stability AI</option>
             </select>
             <div data-sd-source="auto">
                 <label for="sd_auto_url">SD Web UI URL</label>

                 </div>
             </div>
             <div data-sd-source="stability">
-                <label for="sd_stability_key">API Key</label>
+                <div class="flex-container flexnowrap alignItemsBaseline marginBot5">
-                <div class="flex-container flexnowrap">
+                    <strong class="flex1" data-i18n="API Key">API Key</strong>
-                    <input id="sd_stability_key" type="password" class="text_pole flex1" placeholder="Enter your Stability AI API key" />
+                    <div id="sd_stability_key" class="menu_button menu_button_icon">
-                    <div id="sd_stability_validate" class="menu_button menu_button_icon">
+                        <i class="fa-fw fa-solid fa-key"></i>
-                        <i class="fa-solid fa-check"></i>
+                        <span data-i18n="Click to set">Click to set</span>
-                        <span data-i18n="Connect">
-                            Connect
-                        </span>
                     </div>
                 </div>
-                <i>You can find your API key in the Stability AI dashboard.</i>
+                <div class="marginBot5">
+                    <i data-i18n="You can find your API key in the Stability AI dashboard.">
+                        You can find your API key in the Stability AI dashboard.
+                    </i>
+                </div>
                 <div class="flex-container">
                     <div class="flex1">
-                        <label for="sd_stability_engine">Engine</label>
+                        <label for="sd_stability_style_preset" data-i18n="Style Preset">Style Preset</label>
-                        <select id="sd_stability_engine">
-                            <option value="v2beta">V2beta Image Generation</option>
-                        </select>
-                    </div>
-                    <div class="flex1">
-                        <label for="sd_stability_style_preset">Style Preset</label>
                         <select id="sd_stability_style_preset">
                             <option value="anime">Anime</option>
                             <option value="3d-model">3D Model</option>

                         </select>
                     </div>
                 </div>
+            </div>
-                <div class="flex-container">
-                    <div class="flex1">
-                        <label for="sd_stability_aspect_ratio">Aspect Ratio</label>
-                        <select id="sd_stability_aspect_ratio">
-                            <option value="16:9">16:9</option>
-                            <option value="1:1">1:1</option>
-                            <option value="21:9">21:9</option>
-                            <option value="2:3">2:3</option>
-                            <option value="3:2">3:2</option>
-                            <option value="4:5">4:5</option>
-                            <option value="5:4">5:4</option>
-                            <option value="9:16">9:16</option>
-                            <option value="9:21">9:21</option>
-                        </select>
-                    </div>
-                    <div class="flex1">
-                        <label for="sd_stability_seed">Seed</label>
-                        <input id="sd_stability_seed" type="number" class="text_pole" value="0" min="0" max="4294967295" />
-                    </div>
-                </div>
-                <div class="flex-container">
-                    <div class="flex1">
-                        <label for="sd_stability_output_format">Output Format</label>
-                        <select id="sd_stability_output_format">
-                            <option value="png">PNG</option>
-                            <option value="webp">WebP</option>
-                            <option value="jpeg">JPEG</option>
-                        </select>
-                    </div>
-                </div>
-            </div>
             <div class="flex-container">
                 <div class="flex1">
                     <label for="sd_model" data-i18n="Model">Model</label>

                 </label>
             </div>
-            <div data-sd-source="novel,togetherai,pollinations,comfy,drawthings,vlad,auto,horde,extras" class="marginTop5">
+            <div data-sd-source="novel,togetherai,pollinations,comfy,drawthings,vlad,auto,horde,extras,stability" class="marginTop5">
                 <label for="sd_seed">
                     <span data-i18n="Seed">Seed</span>
                     <small data-i18n="(-1 for random)">(-1 for random)</small>

public/scripts/secrets.js+1 -0 Viewed

     FEATHERLESS: 'api_key_featherless',
     ZEROONEAI: 'api_key_01ai',
     HUGGINGFACE: 'api_key_huggingface',
+    STABILITY: 'api_key_stability',
 };
 const INPUT_MAP = {

src/endpoints/secrets.js+1 -0 Viewed

     FEATHERLESS: 'api_key_featherless',
     ZEROONEAI: 'api_key_01ai',
     HUGGINGFACE: 'api_key_huggingface',
+    STABILITY: 'api_key_stability',
 };
 // These are the keys that are safe to expose, even if allowKeysExposure is false

src/endpoints/stable-diffusion.js+60 -0 Viewed

 const writeFileAtomicSync = require('write-file-atomic').sync;
 const { jsonParser } = require('../express-common');
 const { readSecret, SECRET_KEYS } = require('./secrets.js');
+const FormData = require('form-data');
 /**
  * Sanitizes a string.

+    }
 });
+const stability = express.Router();
+stability.post('/generate', jsonParser, async (request, response) => {
+    try {
+        const key = readSecret(request.user.directories, SECRET_KEYS.STABILITY);
+        if (!key) {
+            console.log('Stability AI key not found.');
+            return response.sendStatus(400);
+        }
+        const { payload, model } = request.body;
+        const formData = new FormData();
+        for (const [key, value] of Object.entries(payload)) {
+            if (value !== undefined) {
+                formData.append(key, String(value));
+            }
+        }
+        let apiUrl;
+        switch (model) {
+            case 'stable-image-ultra':
+                apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/ultra';
+                break;
+            case 'stable-image-core':
+                apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/core';
+                break;
+            case 'stable-diffusion-3':
+                apiUrl = 'https://api.stability.ai/v2beta/stable-image/generate/sd3';
+                break;
+            default:
+                throw new Error('Invalid Stability AI model selected');
+        }
+        const result = await fetch(apiUrl, {
+            method: 'POST',
+            headers: {
+                'Authorization': `Bearer ${key}`,
+                'Accept': 'image/*',
+            },
+            body: formData,
+        });
+        if (!result.ok) {
+            const text = await result.text();
+            console.log('Stability AI returned an error.', result.status, result.statusText, text);
+            return response.sendStatus(500);
+        }
+        const buffer = await result.buffer();
+        return response.send(buffer);
+    } catch (error) {
+        console.log(error);
+        return response.sendStatus(500);
+    }
+});
 router.use('/comfy', comfy);
 router.use('/together', together);
 router.use('/drawthings', drawthings);
 router.use('/pollinations', pollinations);
+router.use('/stability', stability);
 module.exports = { router };