Add Stable Horde image gen

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

@@ -3,9 +3,11 @@ import {
     saveSettingsDebounced,
     systemUserName,
     hideSwipeButtons,
-    showSwipeButtons
+    showSwipeButtons,
+    callPopup,
+    getRequestHeaders
 } from "../../../script.js";
-import { getApiUrl, getContext, extension_settings, defaultRequestArgs } from "../../extensions.js";
+import { getApiUrl, getContext, extension_settings, defaultRequestArgs, modules } from "../../extensions.js";
 import { stringFormat, initScrollHeight, resetScrollHeight } from "../../utils.js";
 export { MODULE_NAME };
 
@@ -94,6 +96,9 @@ const defaultSettings = {
     negative_prompt: 'lowres, bad anatomy, bad hands, text, error, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry',
     sampler: 'DDIM',
     model: '',
+
+    horde: false,
+    horde_nsfw: false,
 }
 
 async function loadSettings() {
@@ -107,11 +112,10 @@ async function loadSettings() {
     $('#sd_negative_prompt').val(extension_settings.sd.negative_prompt).trigger('input');
     $('#sd_width').val(extension_settings.sd.width).trigger('input');
     $('#sd_height').val(extension_settings.sd.height).trigger('input');
-
+    $('#sd_horde').prop('checked', extension_settings.sd.horde);
+    $('#sd_horde_nsfw').prop('checked', extension_settings.sd.horde_nsfw);
 
     await Promise.all([loadSamplers(), loadModels()]);
-
-
 }
 
 function onScaleInput() {
@@ -155,10 +159,29 @@ function onHeightInput() {
     saveSettingsDebounced();
 }
 
+async function onHordeInput() {
+    extension_settings.sd.model = null;
+    extension_settings.sd.sampler = null;
+    extension_settings.sd.horde = !!$(this).prop('checked');
+    saveSettingsDebounced();
+    await Promise.all([loadModels(), loadSamplers()]);
+}
+
+async function onHordeNsfwInput() {
+    extension_settings.sd.horde_nsfw = !!$(this).prop('checked');
+    saveSettingsDebounced();
+}
+
 async function onModelChange() {
     extension_settings.sd.model = $('#sd_model').find(':selected').val();
     saveSettingsDebounced();
 
+    if (extension_settings.sd.horde == false) {
+        await updateExtrasRemoteModel();
+    }
+}
+
+async function updateExtrasRemoteModel() {
     const url = new URL(getApiUrl());
     url.pathname = '/api/image/model';
     const getCurrentModelResult = await fetch(url, {
@@ -173,25 +196,86 @@ async function onModelChange() {
 }
 
 async function loadSamplers() {
+    $('#sd_sampler').empty();
+    let samplers = [];
+
+    if (extension_settings.sd.horde) {
+        samplers = await loadHordeSamplers();
+    } else {
+        samplers = await loadExtrasSamplers();
+    }
+
+    for (const sampler of samplers) {
+        const option = document.createElement('option');
+        option.innerText = sampler;
+        option.value = sampler;
+        option.selected = sampler === extension_settings.sd.sampler;
+        $('#sd_sampler').append(option);
+    }
+}
+
+async function loadHordeSamplers() {
+    const result = await fetch('/horde_samplers', {
+        method: 'POST',
+        headers: getRequestHeaders(),
+    });
+
+    if (result.ok) {
+        const data = await result.json();
+        return data;
+    }
+
+    return [];
+}
+
+async function loadExtrasSamplers() {
     const url = new URL(getApiUrl());
     url.pathname = '/api/image/samplers';
     const result = await fetch(url, defaultRequestArgs);
 
     if (result.ok) {
         const data = await result.json();
-        const samplers = data.samplers;
+        return data.samplers;
-
-        for (const sampler of samplers) {
-            const option = document.createElement('option');
-            option.innerText = sampler;
-            option.value = sampler;
-            option.selected = sampler === extension_settings.sd.sampler;
-            $('#sd_sampler').append(option);
-        }
     }
+
+    return [];
 }
 
 async function loadModels() {
+    $('#sd_model').empty();
+    let models = [];
+
+    if (extension_settings.sd.horde) {
+        models = await loadHordeModels();
+    } else {
+        models = await loadExtrasModels();
+    }
+
+    for (const model of models) {
+        const option = document.createElement('option');
+        option.innerText = model.text;
+        option.value = model.value;
+        option.selected = model.value === extension_settings.sd.model;
+        $('#sd_model').append(option);
+    }
+}
+
+async function loadHordeModels() {
+    const result = await fetch('/horde_models', {
+        method: 'POST',
+        headers: getRequestHeaders(),
+    });
+
+    if (result.ok) {
+        const data = await result.json();
+        const models = data.map(x => ({ value: x.name, text: `${x.name} (ETA: ${x.eta}s, Queue: ${x.queued}, Workers: ${x.count})` }));
+        return models;
+    }
+
+    return [];
+}
+
+async function loadExtrasModels() {
     const url = new URL(getApiUrl());
     url.pathname = '/api/image/model';
     const getCurrentModelResult = await fetch(url, defaultRequestArgs);
@@ -206,16 +290,11 @@ async function loadModels() {
 
     if (getModelsResult.ok) {
         const data = await getModelsResult.json();
-        const models = data.models;
+        const view_models = data.models.map(x => ({ value: x, text: x }));
-
+        return view_models;
-        for (const model of models) {
-            const option = document.createElement('option');
-            option.innerText = model;
-            option.value = model;
-            option.selected = model === extension_settings.sd.model;
-            $('#sd_model').append(option);
-        }
     }
+
+    return [];
 }
 
 function getGenerationType(prompt) {
@@ -257,6 +336,14 @@ async function generatePicture(_, trigger) {
         return;
     }
 
+    if (!modules.includes('sd') && !extension_settings.sd.horde) {
+        callPopup("Extensions API is not connected or doesn't provide SD module. Enable Stable Horde to generate images.", 'text');
+        return;
+    }
+
+    extension_settings.sd.sampler = $('#sd_sampler').find(':selected').val();
+    extension_settings.sd.model = $('#sd_model').find(':selected').val();
+
     trigger = trigger.trim();
     const generationMode = getGenerationType(trigger);
     console.log('Generation mode', generationMode, 'triggered with', trigger);
@@ -279,30 +366,10 @@ async function generatePicture(_, trigger) {
 
         console.log('Processed Stable Diffusion prompt:', prompt);
 
-        const url = new URL(getApiUrl());
+        if (extension_settings.sd.horde) {
-        url.pathname = '/api/image';
+            await generateHordeImage(prompt);
-        const result = await fetch(url, {
+        } else {
-            method: 'POST',
+            await generateExtrasImage(prompt);
-            headers: postHeaders,
-            body: JSON.stringify({
-                prompt: prompt,
-                sampler: extension_settings.sd.sampler,
-                steps: extension_settings.sd.steps,
-                scale: extension_settings.sd.scale,
-                width: extension_settings.sd.width,
-                height: extension_settings.sd.height,
-                prompt_prefix: extension_settings.sd.prompt_prefix,
-                negative_prompt: extension_settings.sd.negative_prompt,
-                restore_faces: true,
-                face_restoration_model: 'GFPGAN',
-
-            }),
-        });
-
-        if (result.ok) {
-            const data = await result.json();
-            const base64Image = `data:image/jpeg;base64,${data.image}`;
-            sendMessage(prompt, base64Image);
         }
     } catch (err) {
         console.trace(err);
@@ -314,6 +381,59 @@ async function generatePicture(_, trigger) {
     }
 }
 
+async function generateExtrasImage(prompt) {
+    const url = new URL(getApiUrl());
+    url.pathname = '/api/image';
+    const result = await fetch(url, {
+        method: 'POST',
+        headers: postHeaders,
+        body: JSON.stringify({
+            prompt: prompt,
+            sampler: extension_settings.sd.sampler,
+            steps: extension_settings.sd.steps,
+            scale: extension_settings.sd.scale,
+            width: extension_settings.sd.width,
+            height: extension_settings.sd.height,
+            prompt_prefix: extension_settings.sd.prompt_prefix,
+            negative_prompt: extension_settings.sd.negative_prompt,
+            restore_faces: true,
+            face_restoration_model: 'GFPGAN',
+        }),
+    });
+
+
+    if (result.ok) {
+        const data = await result.json();
+        const base64Image = `data:image/jpeg;base64,${data.image}`;
+        sendMessage(prompt, base64Image);
+    }
+}
+
+async function generateHordeImage(prompt) {
+    const result = await fetch('/horde_generateimage', {
+        method: 'POST',
+        headers: getRequestHeaders(),
+        body: JSON.stringify({
+            prompt: prompt,
+            sampler: extension_settings.sd.sampler,
+            steps: extension_settings.sd.steps,
+            scale: extension_settings.sd.scale,
+            width: extension_settings.sd.width,
+            height: extension_settings.sd.height,
+            prompt_prefix: extension_settings.sd.prompt_prefix,
+            negative_prompt: extension_settings.sd.negative_prompt,
+            model: extension_settings.sd.model,
+            nsfw: extension_settings.sd.horde_nsfw,
+        }),
+    });
+
+    if (result.ok) {
+        const data = await result.text();
+        const base64Image = `data:image/webp;base64,${data}`;
+        sendMessage(prompt, base64Image);
+    }
+}
+
 async function sendMessage(prompt, image) {
     const context = getContext();
     const messageText = `[${context.name2} sends a picture that contains: ${prompt}]`;
@@ -386,16 +506,6 @@ function addSDGenButtons() {
 async function moduleWorker() {
     const context = getContext();
 
-    /*     if (context.onlineStatus === 'no_connection') {
-            $('#sd_gen').hide(200);
-        } else if ($("#send_but").css('display') === 'flex') {
-            $('#sd_gen').show(200);
-            $("#sd_gen_wait").hide(200);
-        } else {
-            $('#sd_gen').hide(200);
-            $("#sd_gen_wait").show(200);
-        } */
-
     context.onlineStatus === 'no_connection'
         ? $('#sd_gen').hide(200)
         : $('#sd_gen').show(200)
@@ -444,6 +554,16 @@ jQuery(async () => {
         </div>
         <div class="inline-drawer-content">
             <small><i>Use slash commands to generate images. Type <span class="monospace">/help</span> in chat for more details</i></small>
+            <br>
+            <small><i>Hint: Save an API key in Horde KoboldAI API settings to use it here.</i></small>
+            <label class="checkbox_label">
+                <input id="sd_horde" type="checkbox" />
+                Use Stable Horde
+            </label>
+            <label style="margin-left:1em;" class="checkbox_label">
+                <input id="sd_horde_nsfw" type="checkbox" />
+                Allow NSFW images
+            </label>
             <label for="sd_scale">CFG Scale (<span id="sd_scale_value"></span>)</label>
             <input id="sd_scale" type="range" min="${defaultSettings.scale_min}" max="${defaultSettings.scale_max}" step="${defaultSettings.scale_step}" value="${defaultSettings.scale}" />
             <label for="sd_steps">Sampling steps (<span id="sd_steps_value"></span>)</label>
@@ -472,6 +592,8 @@ jQuery(async () => {
     $('#sd_negative_prompt').on('input', onNegativePromptInput);
     $('#sd_width').on('input', onWidthInput);
     $('#sd_height').on('input', onHeightInput);
+    $('#sd_horde').on('input', onHordeInput);
+    $('#sd_horde_nsfw').on('input', onHordeNsfwInput);
 
     $('.sd_settings .inline-drawer-toggle').on('click', function () {
         initScrollHeight($("#sd_prompt_prefix"));

public/scripts/extensions/stable-diffusion/style.css

@@ -1,4 +1,4 @@
-.sd_settings label {
+.sd_settings label:not(.checkbox_label) {
     display: block;
 }
 
@@ -16,7 +16,6 @@
     display: flex;
     align-items: center;
     justify-content: center;
-
 }
 
 #sd_gen:hover {

server.js

@@ -2925,10 +2925,21 @@ app.post('/viewsecrets', jsonParser, async (_, response) => {
     }
 });
 
-app.post('/horde_generateimage', async (request, response) => {
+app.post('/horde_samplers', jsonParser, async (_, response) => {
+    const samplers = Object.values(ai_horde.ModelGenerationInputStableSamplers);
+    response.send(samplers);
+});
+
+app.post('/horde_models', jsonParser, async (_, response) => {
+    const models = await ai_horde.getModels();
+    response.send(models);
+});
+
+app.post('/horde_generateimage', jsonParser, async (request, response) => {
     const MAX_ATTEMPTS = 100;
     const CHECK_INTERVAL = 3000;
     const api_key_horde = readSecret(SECRET_KEYS.HORDE) || ANONYMOUS_KEY;
+    console.log('Stable Horde request:', request.body);
     const generation = await ai_horde.postAsyncImageGenerate(
         {
             prompt: `${request.body.prompt_prefix} ${request.body.prompt} ### ${request.body.negative_prompt}`,
@@ -2950,12 +2961,17 @@ app.post('/horde_generateimage', async (request, response) => {
     for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt++) {
         await delay(CHECK_INTERVAL);
         const check = await ai_horde.getImageGenerationCheck(generation.id);
+        console.log(check);
     
         if (check.done) {
             const result = await ai_horde.getImageGenerationStatus(generation.id);
             return response.send(result.generations[0].img);
         }
 
+        if (!check.is_possible) {
+            return response.sendStatus(503);
+        }
+
         if (check.faulted) {
             return response.sendStatus(500);
         }