Merge pull request #2471 from WBlair1/staging Stable image from StabilityAI api

ce71c0ef86837f7223289d36f983caee59a8f64d

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

Signed

Parents: 8fe111cc7155c6d21cd1cb848ebd0b19e13ab675 220051b93375359e39cdb843f426739f5cef20e6

public/scripts/extensions/stable-diffusion/index.js+147 -3 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, writeSecret } from '../../secrets.js';
 import { getNovelUnlimitedImageGeneration, getNovelAnlas, loadNovelSubscriptionData } from '../../nai-settings.js';
 import { getMultimodalCaption } from '../shared.js';
 import { SlashCommandParser } from '../../slash-commands/SlashCommandParser.js';

     togetherai: 'togetherai',
     drawthings: 'drawthings',
     pollinations: 'pollinations',
 +    stability: 'stability',
 };
 const initiators = {

     wand_visible: false,
     command_visible: false,
     interactive_visible: false,
++
 +    // Stability AI settings
 +    stability_style_preset: 'anime',
 };
 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_style_preset').val(extension_settings.sd.stability_style_preset);
     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 String(refinedPrompt);
         } else {
             throw new Error('Generation aborted by user.');
+        }

     extension_settings.sd.comfy_workflow = $('#sd_comfy_workflow').find(':selected').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() {
 +    extension_settings.sd.stability_style_preset = String($('#sd_stability_style_preset').val());
 +    saveSettingsDebounced();
 +}
++
 async function changeComfyWorkflow(_, name) {
     name = name.replace(/(\.json)?$/i, '.json');
     if ($(`#sd_comfy_workflow > [value="${name}"]`).length > 0) {

     extension_settings.sd.model = $('#sd_model').find(':selected').val();
     saveSettingsDebounced();
     const cloudSources = [sources.horde, sources.novel, sources.openai, sources.togetherai, sources.pollinations, sources.stability];
     if (cloudSources.includes(extension_settings.sd.source)) {
         return;

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

         case sources.pollinations:
             models = await loadPollinationsModels();
             break;
 +        case sources.stability:
 +            models = await loadStabilityModels();
 +            break;
+    }
     for (const model of models) {

+    }
+}
 +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' },
 +        { value: 'stable-diffusion-3', text: 'Stable Diffusion 3' },
 +    ];
 +}
++
 async function loadPollinationsModels() {
     return [
+        {

         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.pollinations:
                 result = await generatePollinationsImage(prefixedPrompt, negativePrompt);
                 break;
 +            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) {
 +            throw new Error(`HTTP ${response.status}: ${response.statusText}`);
 +        }
++
 +        const base64Image = await response.text();
++
 +        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.

@@ -3228,6 +3368,8 @@ function isValidState() {
3228	3368	return secret_state[SECRET_KEYS.TOGETHERAI];
3229	3369	case sources.pollinations:
3230	3370	return true;
	3371	+ case sources.stability:
	3372	+ return secret_state[SECRET_KEYS.STABILITY];
3231	3373	}
3232	3374	}
3233	3375

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

public/scripts/extensions/stable-diffusion/settings.html+40 -1 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>

                     </label>
                 </div>
             </div>
 +            <div data-sd-source="stability">
 +                <div class="flex-container flexnowrap alignItemsBaseline marginBot5">
 +                    <strong class="flex1" data-i18n="API Key">API Key</strong>
 +                    <div id="sd_stability_key" class="menu_button menu_button_icon">
 +                        <i class="fa-fw fa-solid fa-key"></i>
 +                        <span data-i18n="Click to set">Click to set</span>
 +                    </div>
 +                </div>
 +                <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_style_preset" data-i18n="Style Preset">Style Preset</label>
 +                        <select id="sd_stability_style_preset">
 +                            <option value="anime">Anime</option>
 +                            <option value="3d-model">3D Model</option>
 +                            <option value="analog-film">Analog Film</option>
 +                            <option value="cinematic">Cinematic</option>
 +                            <option value="comic-book">Comic Book</option>
 +                            <option value="digital-art">Digital Art</option>
 +                            <option value="enhance">Enhance</option>
 +                            <option value="fantasy-art">Fantasy Art</option>
 +                            <option value="isometric">Isometric</option>
 +                            <option value="line-art">Line Art</option>
 +                            <option value="low-poly">Low Poly</option>
 +                            <option value="modeling-compound">Modeling Compound</option>
 +                            <option value="neon-punk">Neon Punk</option>
 +                            <option value="origami">Origami</option>
 +                            <option value="photographic">Photographic</option>
 +                            <option value="pixel-art">Pixel Art</option>
 +                            <option value="tile-texture">Tile Texture</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,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+61 -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,
 +            timeout: 0,
 +        });
++
 +        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.toString('base64'));
 +    } 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 };