Move Stability generation to backend

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

@@ -22,7 +22,7 @@ import { getApiUrl, getContext, extension_settings, doExtrasFetch, modules, rend
 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';
@@ -285,11 +285,7 @@ const defaultSettings = {
     interactive_visible: false,
 
     // Stability AI settings
-    stability_api_key: '',
-    stability_engine: 'V2beta Image Generation',
-    stability_style_preset: "anime",
-    stability_aspect_ratio: '1:1',
-    stability_output_format: 'png',
+    stability_style_preset: 'anime',
 };
 
 const writePromptFieldsDebounced = debounce(writePromptFields, debounce_timeout.relaxed);
@@ -452,11 +448,7 @@ async function loadSettings() {
     $('#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');
@@ -684,7 +676,7 @@ async function refinePrompt(prompt, allowExpand, isNegative = false) {
         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.');
         }
@@ -1097,14 +1089,19 @@ function onComfyWorkflowChange() {
     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() {
-    extension_settings.sd.stability_engine = $('#sd_stability_engine').val();
-    saveSettingsDebounced();
+async function onStabilityKeyClick() {
+    const popupText = 'Stability AI API Key:';
+    const key = await callGenericPopup(popupText, POPUP_TYPE.INPUT);
+
+    if (!key) {
+        return;
+    }
+
+    await writeSecret(SECRET_KEYS.STABILITY, String(key));
+
+    toastr.success('API Key saved');
+    await loadSettingOptions();
 }
 
 function onStabilityStylePresetChange() {
@@ -1112,17 +1109,6 @@ function onStabilityStylePresetChange() {
     saveSettingsDebounced();
 }
 
-function onStabilityAspectRatioChange() {
-    extension_settings.sd.stability_aspect_ratio = $('#sd_stability_aspect_ratio').val();
-    saveSettingsDebounced();
-}
-
-function onStabilityOutputFormatChange() {
-    extension_settings.sd.stability_output_format = $('#sd_stability_output_format').val();
-    saveSettingsDebounced();
-}
-
-
 async function changeComfyWorkflow(_, name) {
     name = name.replace(/(\.json)?$/i, '.json');
     if ($(`#sd_comfy_workflow > [value="${name}"]`).length > 0) {
@@ -1441,6 +1427,9 @@ async function loadSamplers() {
         case sources.pollinations:
             samplers = ['N/A'];
             break;
+        case sources.stability:
+            samplers = ['N/A'];
+            break;
     }
 
     for (const sampler of samplers) {
@@ -1643,80 +1632,9 @@ async function loadModels() {
     }
 }
 
-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' },
@@ -2055,6 +1973,9 @@ async function loadSchedulers() {
         case sources.comfy:
             schedulers = await loadComfySchedulers();
             break;
+        case sources.stability:
+            schedulers = ['N/A'];
+            break;
     }
 
     for (const scheduler of schedulers) {
@@ -2128,6 +2049,9 @@ async function loadVaes() {
         case sources.comfy:
             vaes = await loadComfyVaes();
             break;
+        case sources.stability:
+            vaes = ['N/A'];
+            break;
     }
 
     for (const vae of vaes) {
@@ -2611,7 +2535,6 @@ async function sendGenerationRequest(generationType, prompt, additionalNegativeP
             case sources.stability:
                 result = await generateStabilityImage(prefixedPrompt, negativePrompt);
                 break;
-    
         }
 
         if (!result.data) {
@@ -2635,6 +2558,12 @@ async function sendGenerationRequest(generationType, prompt, additionalNegativeP
     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',
@@ -2659,6 +2588,12 @@ async function generateTogetherAIImage(prompt, negativePrompt) {
     }
 }
 
+/**
+ * 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',
@@ -2727,6 +2662,86 @@ async function generateExtrasImage(prompt, negativePrompt) {
     }
 }
 
+/**
+ * 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.
  *
@@ -3356,7 +3371,7 @@ function isValidState() {
         case sources.pollinations:
             return true;
         case sources.stability:
-            return !!extension_settings.sd.stability_key;
+            return secret_state[SECRET_KEYS.STABILITY];
     }
 }
 
@@ -3584,12 +3599,9 @@ jQuery(async () => {
     $('#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_engine').on('change', onStabilityEngineChange);
+    $('#sd_stability_key').on('click', onStabilityKeyClick);
     $('#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'));