From e3e6fa22187f5e25a661770a385142f747292f05 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Sun, 3 Sep 2023 00:41:26 +0300
Subject: [PATCH] Connect to AUTO111 without Extras. Add NAI Diffusion.

---
 .../extensions/stable-diffusion/index.js      | 501 +++++++++++++-----
 .../extensions/stable-diffusion/settings.html |  94 ++++
 server.js                                     | 234 +++++++-
 3 files changed, 692 insertions(+), 137 deletions(-)
 create mode 100644 public/scripts/extensions/stable-diffusion/settings.html

diff --git a/public/scripts/extensions/stable-diffusion/index.js b/public/scripts/extensions/stable-diffusion/index.js
index 38d26cb4d..efb190770 100644
--- a/public/scripts/extensions/stable-diffusion/index.js
+++ b/public/scripts/extensions/stable-diffusion/index.js
@@ -10,11 +10,13 @@ import {
     appendImageToMessage,
     generateQuietPrompt,
     this_chid,
+    getCurrentChatId,
 } from "../../../script.js";
-import { getApiUrl, getContext, extension_settings, doExtrasFetch, modules } from "../../extensions.js";
+import { getApiUrl, getContext, extension_settings, doExtrasFetch, modules, renderExtensionTemplate } from "../../extensions.js";
 import { selected_group } from "../../group-chats.js";
 import { stringFormat, initScrollHeight, resetScrollHeight, timestampToMoment, getCharaFilename, saveBase64AsFile } from "../../utils.js";
 import { getMessageTimeStamp, humanizedDateTime } from "../../RossAscends-mods.js";
+import { SECRET_KEYS, secret_state } from "../../secrets.js";
 export { MODULE_NAME };
 
 // Wraps a string into monospace font-face span
@@ -27,10 +29,12 @@ const p = a => `<p>${a}</p>`
 const MODULE_NAME = 'sd';
 const UPDATE_INTERVAL = 1000;
 
-const postHeaders = {
-    'Content-Type': 'application/json',
-    'Bypass-Tunnel-Reminder': 'bypass',
-};
+const sources = {
+    extras: 'extras',
+    horde: 'horde',
+    auto: 'auto',
+    novel: 'novel',
+}
 
 const generationMode = {
     CHARACTER: 0,
@@ -116,6 +120,8 @@ const helpString = [
 ].join('<br>');
 
 const defaultSettings = {
+    source: sources.extras,
+
     // CFG Scale
     scale_min: 1,
     scale_max: 30,
@@ -153,13 +159,31 @@ const defaultSettings = {
     refine_mode: false,
 
     prompts: promptTemplates,
+
+    // AUTOMATIC1111 settings
+    auto_url: 'http://localhost:7860',
+}
+
+function toggleSourceControls() {
+    $('.sd_settings [data-sd-source]').each(function () {
+        const source = $(this).data('sd-source');
+        $(this).toggle(source === extension_settings.sd.source);
+    });
 }
 
 async function loadSettings() {
+    // Initialize settings
     if (Object.keys(extension_settings.sd).length === 0) {
         Object.assign(extension_settings.sd, defaultSettings);
     }
 
+    // Insert missing settings
+    for (const [key, value] of Object.entries(defaultSettings)) {
+        if (extension_settings.sd[key] === undefined) {
+            extension_settings.sd[key] = value;
+        }
+    }
+
     if (extension_settings.sd.prompts === undefined) {
         extension_settings.sd.prompts = promptTemplates;
     }
@@ -175,6 +199,7 @@ async function loadSettings() {
         extension_settings.sd.character_prompts = {};
     }
 
+    $('#sd_source').val(extension_settings.sd.source);
     $('#sd_scale').val(extension_settings.sd.scale).trigger('input');
     $('#sd_steps').val(extension_settings.sd.steps).trigger('input');
     $('#sd_prompt_prefix').val(extension_settings.sd.prompt_prefix).trigger('input');
@@ -187,7 +212,9 @@ async function loadSettings() {
     $('#sd_restore_faces').prop('checked', extension_settings.sd.restore_faces);
     $('#sd_enable_hr').prop('checked', extension_settings.sd.enable_hr);
     $('#sd_refine_mode').prop('checked', extension_settings.sd.refine_mode);
+    $('#sd_auto_url').val(extension_settings.sd.auto_url);
 
+    toggleSourceControls();
     addPromptTemplates();
 
     await Promise.all([loadSamplers(), loadModels()]);
@@ -332,10 +359,11 @@ function onHeightInput() {
     saveSettingsDebounced();
 }
 
-async function onHordeInput() {
+async function onSourceChange() {
+    extension_settings.sd.source = $('#sd_source').find(':selected').val();
     extension_settings.sd.model = null;
     extension_settings.sd.sampler = null;
-    extension_settings.sd.horde = !!$(this).prop('checked');
+    toggleSourceControls();
     saveSettingsDebounced();
     await Promise.all([loadModels(), loadSamplers()]);
 }
@@ -360,13 +388,92 @@ function onHighResFixInput() {
     saveSettingsDebounced();
 }
 
+function onAutoUrlInput() {
+    extension_settings.sd.auto_url = $('#sd_auto_url').val();
+    saveSettingsDebounced();
+}
+
+async function validateAutoUrl() {
+    try {
+        if (!extension_settings.sd.auto_url) {
+            throw new Error('URL is not set.');
+        }
+
+        const result = await fetch('/api/sd/ping', {
+            method: 'POST',
+            headers: getRequestHeaders(),
+            body: JSON.stringify({ url: extension_settings.sd.auto_url }),
+        });
+
+        if (!result.ok) {
+            throw new Error('SD WebUI returned an error.');
+        }
+
+        await loadSamplers();
+        await loadModels();
+        toastr.success('SD WebUI API connected.');
+    } catch (error) {
+        toastr.error(`Could not validate SD WebUI API: ${error.message}`);
+    }
+}
+
 async function onModelChange() {
     extension_settings.sd.model = $('#sd_model').find(':selected').val();
     saveSettingsDebounced();
 
-    if (!extension_settings.sd.horde) {
+    const cloudSources = [sources.horde, sources.novel];
+
+    if (cloudSources.includes(extension_settings.sd.source)) {
+        return;
+    }
+
+    toastr.info('Updating remote model...', 'Please wait');
+    if (extension_settings.sd.source === sources.extras) {
         await updateExtrasRemoteModel();
     }
+    if (extension_settings.sd.source === sources.auto) {
+        await updateAutoRemoteModel();
+    }
+    toastr.success('Model successfully loaded!', 'Stable Diffusion');
+}
+
+async function getAutoRemoteModel() {
+    try {
+        const result = await fetch('/api/sd/get-model', {
+            method: 'POST',
+            headers: getRequestHeaders(),
+            body: JSON.stringify({ url: extension_settings.sd.auto_url }),
+        });
+
+        if (!result.ok) {
+            throw new Error('SD WebUI returned an error.');
+        }
+
+        const data = await result.text();
+        return data;
+    } catch (error) {
+        console.error(error);
+        return null;
+    }
+}
+
+async function updateAutoRemoteModel() {
+    try {
+        const result = await fetch('/api/sd/set-model', {
+            method: 'POST',
+            headers: getRequestHeaders(),
+            body: JSON.stringify({ url: extension_settings.sd.auto_url, model: extension_settings.sd.model }),
+        });
+
+        if (!result.ok) {
+            throw new Error('SD WebUI returned an error.');
+        }
+
+        console.log('Model successfully updated on SD WebUI remote.');
+    } catch (error) {
+        console.error(error);
+        toastr.error(`Could not update SD WebUI model: ${error.message}`);
+    }
 }
 
 async function updateExtrasRemoteModel() {
@@ -374,7 +481,6 @@ async function updateExtrasRemoteModel() {
     url.pathname = '/api/image/model';
     const getCurrentModelResult = await doExtrasFetch(url, {
         method: 'POST',
-        headers: postHeaders,
         body: JSON.stringify({ model: extension_settings.sd.model }),
     });
 
@@ -387,10 +493,19 @@ async function loadSamplers() {
     $('#sd_sampler').empty();
     let samplers = [];
 
-    if (extension_settings.sd.horde) {
-        samplers = await loadHordeSamplers();
-    } else {
-        samplers = await loadExtrasSamplers();
+    switch (extension_settings.sd.source) {
+        case sources.extras:
+            samplers = await loadExtrasSamplers();
+            break;
+        case sources.horde:
+            samplers = await loadHordeSamplers();
+            break;
+        case sources.auto:
+            samplers = await loadAutoSamplers();
+            break;
+        case sources.novel:
+            samplers = await loadNovelSamplers();
+            break;
     }
 
     for (const sampler of samplers) {
@@ -433,14 +548,63 @@ async function loadExtrasSamplers() {
     return [];
 }
 
+async function loadAutoSamplers() {
+    if (!extension_settings.sd.auto_url) {
+        return [];
+    }
+
+    try {
+        const result = await fetch('/api/sd/samplers', {
+            method: 'POST',
+            headers: getRequestHeaders(),
+            body: JSON.stringify({ url: extension_settings.sd.auto_url }),
+        });
+
+        if (!result.ok) {
+            throw new Error('SD WebUI returned an error.');
+        }
+
+        const data = await result.json();
+        return data;
+    } catch (error) {
+        return [];
+    }
+}
+
+async function loadNovelSamplers() {
+    if (!secret_state[SECRET_KEYS.NOVEL]) {
+        toastr.warning('NovelAI API key is not set.');
+        return [];
+    }
+
+    return [
+        'k_dpmpp_2m',
+        'k_dpmpp_sde',
+        'k_dpmpp_2s_ancestral',
+        'k_euler',
+        'k_euler_ancestral',
+        'k_dpm_fast',
+        'ddim',
+    ];
+}
+
 async function loadModels() {
     $('#sd_model').empty();
     let models = [];
 
-    if (extension_settings.sd.horde) {
-        models = await loadHordeModels();
-    } else {
-        models = await loadExtrasModels();
+    switch (extension_settings.sd.source) {
+        case sources.extras:
+            models = await loadExtrasModels();
+            break;
+        case sources.horde:
+            models = await loadHordeModels();
+            break;
+        case sources.auto:
+            models = await loadAutoModels();
+            break;
+        case sources.novel:
+            models = await loadNovelModels();
+            break;
     }
 
     for (const model of models) {
@@ -495,6 +659,57 @@ async function loadExtrasModels() {
     return [];
 }
 
+async function loadAutoModels() {
+    if (!extension_settings.sd.auto_url) {
+        return [];
+    }
+
+    try {
+        const currentModel = await getAutoRemoteModel();
+
+        if (currentModel) {
+            extension_settings.sd.model = currentModel;
+        }
+
+        const result = await fetch('/api/sd/models', {
+            method: 'POST',
+            headers: getRequestHeaders(),
+            body: JSON.stringify({ url: extension_settings.sd.auto_url }),
+        });
+
+        if (!result.ok) {
+            throw new Error('SD WebUI returned an error.');
+        }
+
+        const data = await result.json();
+        return data;
+    } catch (error) {
+        return [];
+    }
+}
+
+async function loadNovelModels() {
+    if (!secret_state[SECRET_KEYS.NOVEL]) {
+        toastr.warning('NovelAI API key is not set.');
+        return [];
+    }
+
+    return [
+        {
+            value: 'nai-diffusion',
+            text: 'Full',
+        },
+        {
+            value: 'safe-diffusion',
+            text: 'Safe',
+        },
+        {
+            value: 'nai-diffusion-furry',
+            text: 'Furry',
+        },
+    ];
+}
+
 function getGenerationType(prompt) {
     for (const [key, values] of Object.entries(triggerWords)) {
         for (const value of values) {
@@ -537,7 +752,6 @@ function processReply(str) {
     return str;
 }
 
-
 function getRawLastMessage() {
     const getLastUsableMessage = () => {
         for (const message of context.chat.slice().reverse()) {
@@ -565,7 +779,7 @@ async function generatePicture(_, trigger, message, callback) {
         return;
     }
 
-    if (!modules.includes('sd') && !extension_settings.sd.horde) {
+    if (!isValidState()) {
         toastr.warning("Extensions API is not connected or doesn't provide SD module. Enable Stable Horde to generate images.");
         return;
     }
@@ -602,8 +816,8 @@ async function generatePicture(_, trigger, message, callback) {
             const imagePath = base64Image;
             const imgUrl = `url('${encodeURIComponent(base64Image)}')`;
 
-            if ('forceSetBackground' in window) {
-                forceSetBackground(imgUrl);
+            if (typeof window['forceSetBackground'] === 'function') {
+                window['forceSetBackground'](imgUrl);
             } else {
                 toastr.info('Background image will not be preserved.', '"Chat backgrounds" extension is disabled.');
                 $('#bg_custom').css('background-image', imgUrl);
@@ -669,32 +883,57 @@ async function sendGenerationRequest(generationType, prompt, characterName = nul
         ? combinePrefixes(extension_settings.sd.prompt_prefix, getCharacterPrefix())
         : extension_settings.sd.prompt_prefix;
 
-    if (extension_settings.sd.horde) {
-        await generateHordeImage(prompt, prefix, characterName, callback);
-    } else {
-        await generateExtrasImage(prompt, prefix, characterName, callback);
+    const prefixedPrompt = combinePrefixes(prefix, prompt);
+
+    let result = { format: '', data: '' };
+    const currentChatId = getCurrentChatId();
+
+    try {
+        switch (extension_settings.sd.source) {
+            case sources.extras:
+                result = await generateExtrasImage(prefixedPrompt);
+                break;
+            case sources.horde:
+                result = await generateHordeImage(prefixedPrompt);
+                break;
+            case sources.auto:
+                result = await generateAutoImage(prefixedPrompt);
+                break;
+            case sources.novel:
+                result = await generateNovelImage(prefixedPrompt);
+                break;
+        }
+
+        if (!result.data) {
+            throw new Error();
+        }
+    } catch (err) {
+        toastr.error('Image generation failed. Please try again', 'Stable Diffusion');
+        return;
     }
+
+    if (currentChatId !== getCurrentChatId()) {
+        console.warn('Chat changed, aborting SD result saving');
+        toastr.warning('Chat changed, generated image discarded.', 'Stable Diffusion');
+        return;
+    }
+
+    const filename = `${characterName}_${humanizedDateTime()}`;
+    const base64Image = await saveBase64AsFile(result.data, characterName, filename, result.format);
+    callback ? callback(prompt, base64Image) : sendMessage(prompt, base64Image);
 }
 
 /**
- * Generates an "extras" image using a provided prompt and other settings,
- * then saves the generated image and either invokes a callback or sends a message with the image.
+ * Generates an "extras" image using a provided prompt and other settings.
  *
  * @param {string} prompt - The main instruction used to guide the image generation.
- * @param {string} prefix - Additional context or prefix to guide the image generation.
- * @param {string} characterName - The name used to determine the sub-directory for saving.
- * @param {function} [callback] - Optional callback function invoked with the prompt and saved image.
- *                                If not provided, `sendMessage` is called instead.
- *
- * @returns {Promise<void>} - A promise that resolves when the image generation and processing are complete.
+ * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete.
  */
-async function generateExtrasImage(prompt, prefix, characterName, callback) {
-    console.debug(extension_settings.sd);
+async function generateExtrasImage(prompt) {
     const url = new URL(getApiUrl());
     url.pathname = '/api/image';
     const result = await doExtrasFetch(url, {
         method: 'POST',
-        headers: postHeaders,
         body: JSON.stringify({
             prompt: prompt,
             sampler: extension_settings.sd.sampler,
@@ -702,7 +941,6 @@ async function generateExtrasImage(prompt, prefix, characterName, callback) {
             scale: extension_settings.sd.scale,
             width: extension_settings.sd.width,
             height: extension_settings.sd.height,
-            prompt_prefix: prefix,
             negative_prompt: extension_settings.sd.negative_prompt,
             restore_faces: !!extension_settings.sd.restore_faces,
             enable_hr: !!extension_settings.sd.enable_hr,
@@ -712,28 +950,19 @@ async function generateExtrasImage(prompt, prefix, characterName, callback) {
 
     if (result.ok) {
         const data = await result.json();
-        //filename should be character name + human readable timestamp + generation mode
-        const filename = `${characterName}_${humanizedDateTime()}`;
-        const base64Image = await saveBase64AsFile(data.image, characterName, filename, "jpg");
-        callback ? callback(prompt, base64Image) : sendMessage(prompt, base64Image);
+        return { format: 'jpg', data: data.image };
     } else {
-        callPopup('Image generation has failed. Please try again.', 'text');
+        throw new Error();
     }
 }
 
 /**
- * Generates a "horde" image using the provided prompt and configuration settings,
- * then saves the generated image and either invokes a callback or sends a message with the image.
+ * Generates a "horde" image using the provided prompt and configuration settings.
  *
  * @param {string} prompt - The main instruction used to guide the image generation.
- * @param {string} prefix - Additional context or prefix to guide the image generation.
- * @param {string} characterName - The name used to determine the sub-directory for saving.
- * @param {function} [callback] - Optional callback function invoked with the prompt and saved image.
- *                                If not provided, `sendMessage` is called instead.
- *
- * @returns {Promise<void>} - A promise that resolves when the image generation and processing are complete.
+ * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete.
  */
-async function generateHordeImage(prompt, prefix, characterName, callback) {
+async function generateHordeImage(prompt) {
     const result = await fetch('/horde_generateimage', {
         method: 'POST',
         headers: getRequestHeaders(),
@@ -744,7 +973,6 @@ async function generateHordeImage(prompt, prefix, characterName, callback) {
             scale: extension_settings.sd.scale,
             width: extension_settings.sd.width,
             height: extension_settings.sd.height,
-            prompt_prefix: prefix,
             negative_prompt: extension_settings.sd.negative_prompt,
             model: extension_settings.sd.model,
             nsfw: extension_settings.sd.horde_nsfw,
@@ -755,11 +983,76 @@ async function generateHordeImage(prompt, prefix, characterName, callback) {
 
     if (result.ok) {
         const data = await result.text();
-        const filename = `${characterName}_${humanizedDateTime()}`;
-        const base64Image = await saveBase64AsFile(data, characterName, filename, "webp");
-        callback ? callback(prompt, base64Image) : sendMessage(prompt, base64Image);
+        return { format: 'webp', data: data };
     } else {
-        toastr.error('Image generation has failed. Please try again.');
+        throw new Error();
+    }
+}
+
+/**
+ * Generates an image in SD WebUI API using the provided prompt and configuration settings.
+ *
+ * @param {string} prompt - The main instruction used to guide the image generation.
+ * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete.
+ */
+async function generateAutoImage(prompt) {
+    const result = await fetch('/api/sd/generate', {
+        method: 'POST',
+        headers: getRequestHeaders(),
+        body: JSON.stringify({
+            url: extension_settings.sd.auto_url,
+            prompt: prompt,
+            negative_prompt: extension_settings.sd.negative_prompt,
+            sampler_name: extension_settings.sd.sampler,
+            steps: extension_settings.sd.steps,
+            cfg_scale: extension_settings.sd.scale,
+            width: extension_settings.sd.width,
+            height: extension_settings.sd.height,
+            enable_hr: !!extension_settings.sd.enable_hr,
+            restore_faces: !!extension_settings.sd.restore_faces,
+            // Ensure generated img is saved to disk
+            save_images: true,
+            send_images: true,
+            do_not_save_grid: false,
+            do_not_save_samples: false,
+        }),
+    });
+
+    if (result.ok) {
+        const data = await result.json();
+        return { format: 'png', data: data.images[0] };
+    } else {
+        throw new Error();
+    }
+}
+
+/**
+ * Generates an image in NovelAI API using the provided prompt and configuration settings.
+ *
+ * @param {string} prompt - The main instruction used to guide the image generation.
+ * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete.
+ */
+async function generateNovelImage(prompt) {
+    const result = await fetch('/api/novelai/generate-image', {
+        method: 'POST',
+        headers: getRequestHeaders(),
+        body: JSON.stringify({
+            prompt: prompt,
+            model: extension_settings.sd.model,
+            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,
+            negative_prompt: extension_settings.sd.negative_prompt,
+        }),
+    });
+
+    if (result.ok) {
+        const data = await result.text();
+        return { format: 'png', data: data };
+    } else {
+        throw new Error();
     }
 }
 
@@ -842,12 +1135,21 @@ function addSDGenButtons() {
     });
 }
 
-function isConnectedToExtras() {
-    return modules.includes('sd');
+function isValidState() {
+    switch (extension_settings.sd.source) {
+        case sources.extras:
+            return modules.includes('sd');
+        case sources.horde:
+            return true;
+        case sources.auto:
+            return !!extension_settings.sd.auto_url;
+        case sources.novel:
+            return secret_state[SECRET_KEYS.NOVEL];
+    }
 }
 
 async function moduleWorker() {
-    if (isConnectedToExtras() || extension_settings.sd.horde) {
+    if (isValidState()) {
         $('#sd_gen').show();
         $('.sd_message_gen').show();
     }
@@ -942,82 +1244,8 @@ $("#sd_dropdown [id]").on("click", function () {
 jQuery(async () => {
     getContext().registerSlashCommand('sd', generatePicture, [], helpString, true, true);
 
-    const settingsHtml = `
-    <div class="sd_settings">
-        <div class="inline-drawer">
-            <div class="inline-drawer-toggle inline-drawer-header">
-                <b>Stable Diffusion</b>
-                <div class="inline-drawer-icon fa-solid fa-circle-chevron-down down"></div>
-            </div>
-            <div class="inline-drawer-content">
-                <small><i>Use slash commands or the bottom Paintbrush button 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 for="sd_refine_mode" class="checkbox_label" title="Allow to edit prompts manually before sending them to generation API">
-                    <input id="sd_refine_mode" type="checkbox" />
-                    Edit prompts before generation
-                </label>
-                <div class="flex-container flexGap5 marginTop10 margin-bot-10px">
-                    <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 from Horde
-                    </label>
-                </div>
-                <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>
-                <input id="sd_steps" type="range" min="${defaultSettings.steps_min}" max="${defaultSettings.steps_max}" step="${defaultSettings.steps_step}" value="${defaultSettings.steps}" />
-                <label for="sd_width">Width (<span id="sd_width_value"></span>)</label>
-                <input id="sd_width" type="range" max="${defaultSettings.dimension_max}" min="${defaultSettings.dimension_min}" step="${defaultSettings.dimension_step}" value="${defaultSettings.width}" />
-                <label for="sd_height">Height (<span id="sd_height_value"></span>)</label>
-                <input id="sd_height" type="range" max="${defaultSettings.dimension_max}" min="${defaultSettings.dimension_min}" step="${defaultSettings.dimension_step}" value="${defaultSettings.height}" />
-                <div><small>Only for Horde or remote Stable Diffusion Web UI:</small></div>
-                <div class="flex-container marginTop10 margin-bot-10px">
-                    <label class="flex1 checkbox_label">
-                        <input id="sd_restore_faces" type="checkbox" />
-                        Restore Faces
-                    </label>
-                    <label class="flex1 checkbox_label">
-                        <input id="sd_enable_hr" type="checkbox" />
-                        Hires. Fix
-                    </label>
-                </div>
-                <label for="sd_model">Stable Diffusion model</label>
-                <select id="sd_model"></select>
-                <label for="sd_sampler">Sampling method</label>
-                <select id="sd_sampler"></select>
-                <div class="flex-container flexGap5 margin-bot-10px">
-                    <label class="checkbox_label">
-                        <input id="sd_horde_karras" type="checkbox" />
-                        Karras (only for Horde, not all samplers supported)
-                    </label>
-                </div>
-                <label for="sd_prompt_prefix">Common prompt prefix</label>
-                <textarea id="sd_prompt_prefix" class="text_pole textarea_compact" rows="3"></textarea>
-                <div id="sd_character_prompt_block">
-                    <label for="sd_character_prompt">Character-specific prompt prefix</label>
-                    <small>Won't be used in groups.</small>
-                    <textarea id="sd_character_prompt" class="text_pole textarea_compact" rows="3" placeholder="Any characteristics that describe the currently selected character. Will be added after a common prefix.&#10;Example: female, green eyes, brown hair, pink shirt"></textarea>
-                </div>
-                <label for="sd_negative_prompt">Negative prompt</label>
-                <textarea id="sd_negative_prompt" class="text_pole textarea_compact" rows="3"></textarea>
-            </div>
-        </div>
-        <div class="inline-drawer">
-            <div class="inline-drawer-toggle inline-drawer-header">
-                <b>SD Prompt Templates</b>
-                <div class="inline-drawer-icon fa-solid fa-circle-chevron-down down"></div>
-            </div>
-            <div id="sd_prompt_templates" class="inline-drawer-content">
-            </div>
-        </div>
-    </div>`;
-
-    $('#extensions_settings').append(settingsHtml);
+    $('#extensions_settings').append(renderExtensionTemplate('stable-diffusion', 'settings', defaultSettings));
+    $('#sd_source').on('change', onSourceChange);
     $('#sd_scale').on('input', onScaleInput);
     $('#sd_steps').on('input', onStepsInput);
     $('#sd_model').on('change', onModelChange);
@@ -1026,13 +1254,14 @@ 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_horde_karras').on('input', onHordeKarrasInput);
     $('#sd_restore_faces').on('input', onRestoreFacesInput);
     $('#sd_enable_hr').on('input', onHighResFixInput);
     $('#sd_refine_mode').on('input', onRefineModeInput);
     $('#sd_character_prompt').on('input', onCharacterPromptInput);
+    $('#sd_auto_validate').on('click', validateAutoUrl);
+    $('#sd_auto_url').on('input', onAutoUrlInput);
     $('#sd_character_prompt_block').hide();
 
     $('.sd_settings .inline-drawer-toggle').on('click', function () {
diff --git a/public/scripts/extensions/stable-diffusion/settings.html b/public/scripts/extensions/stable-diffusion/settings.html
new file mode 100644
index 000000000..110467d37
--- /dev/null
+++ b/public/scripts/extensions/stable-diffusion/settings.html
@@ -0,0 +1,94 @@
+
+<div class="sd_settings">
+    <div class="inline-drawer">
+        <div class="inline-drawer-toggle inline-drawer-header">
+            <b>Stable Diffusion</b>
+            <div class="inline-drawer-icon fa-solid fa-circle-chevron-down down"></div>
+        </div>
+        <div class="inline-drawer-content">
+            <small><i>Use slash commands or the bottom Paintbrush button to generate images. Type <span class="monospace">/help</span> in chat for more details</i></small>
+            <br>
+            <label for="sd_refine_mode" class="checkbox_label" title="Allow to edit prompts manually before sending them to generation API">
+                <input id="sd_refine_mode" type="checkbox" />
+                Edit prompts before generation
+            </label>
+            <label for="sd_source">Source</label>
+            <select id="sd_source">
+                <option value="extras">Extras API (local / remote)</option>
+                <option value="horde">Stable Horde</option>
+                <option value="auto">Stable Diffusion Web UI (AUTOMATIC1111)</option>
+                <option value="novel">NovelAI Diffusion</option>
+            </select>
+            <div data-sd-source="auto">
+                <label for="sd_auto_url">SD Web UI URL</label>
+                <div class="flex-container flexnowrap">
+                    <input id="sd_auto_url" type="text" class="text_pole" placeholder="{{auto_url}}" value="{{auto_url}}" />
+                    <div id="sd_auto_validate" class="menu_button menu_button_icon">
+                        <i class="fa-solid fa-check"></i>
+                        <span data-i18n="Connect">
+                            Connect
+                        </span>
+                    </div>
+                </div>
+                <i><b>Important:</b> run SD Web UI with the <tt>--api</tt> flag! The server must be accessible from the SillyTavern host machine.</i>
+            </div>
+            <div data-sd-source="horde">
+                <i>Hint: Save an API key in Horde KoboldAI API settings to use it here.</i>
+                <label for="sd_horde_nsfw" class="checkbox_label">
+                    <input id="sd_horde_nsfw" type="checkbox" />
+                    <span data-i18n="Allow NSFW images from Horde">
+                        Allow NSFW images from Horde
+                    </span>
+                </label>
+                <label for="sd_horde_karras" class="checkbox_label">
+                    <input id="sd_horde_karras" type="checkbox" />
+                    <span data-i18n="Karras (not all samplers supported)">
+                        Karras (not all samplers supported)
+                    </span>
+                </label>
+            </div>
+            <div data-sd-source="novel">
+                <i>Hint: Save an API key in the NovelAI API settings to use it here.</i>
+            </div>
+            <label for="sd_scale">CFG Scale (<span id="sd_scale_value"></span>)</label>
+            <input id="sd_scale" type="range" min="{{scale_min}}" max="{{scale_max}}" step="{{scale_step}}" value="{{scale}}" />
+            <label for="sd_steps">Sampling steps (<span id="sd_steps_value"></span>)</label>
+            <input id="sd_steps" type="range" min="{{steps_min}}" max="{{steps_max}}" step="{{steps_step}}" value="{{steps}}" />
+            <label for="sd_width">Width (<span id="sd_width_value"></span>)</label>
+            <input id="sd_width" type="range" max="{{dimension_max}}" min="{{dimension_min}}" step="{{dimension_step}}" value="{{width}}" />
+            <label for="sd_height">Height (<span id="sd_height_value"></span>)</label>
+            <input id="sd_height" type="range" max="{{dimension_max}}" min="{{dimension_min}}" step="{{dimension_step}}" value="{{height}}" />
+            <div class="flex-container marginTop10 margin-bot-10px">
+                <label class="flex1 checkbox_label">
+                    <input id="sd_restore_faces" type="checkbox" />
+                    Restore Faces
+                </label>
+                <label class="flex1 checkbox_label">
+                    <input id="sd_enable_hr" type="checkbox" />
+                    Hires. Fix
+                </label>
+            </div>
+            <label for="sd_model">Stable Diffusion model</label>
+            <select id="sd_model"></select>
+            <label for="sd_sampler">Sampling method</label>
+            <select id="sd_sampler"></select>
+            <label for="sd_prompt_prefix">Common prompt prefix</label>
+            <textarea id="sd_prompt_prefix" class="text_pole textarea_compact" rows="3"></textarea>
+            <div id="sd_character_prompt_block">
+                <label for="sd_character_prompt">Character-specific prompt prefix</label>
+                <small>Won't be used in groups.</small>
+                <textarea id="sd_character_prompt" class="text_pole textarea_compact" rows="3" placeholder="Any characteristics that describe the currently selected character. Will be added after a common prefix.&#10;Example: female, green eyes, brown hair, pink shirt"></textarea>
+            </div>
+            <label for="sd_negative_prompt">Negative prompt</label>
+            <textarea id="sd_negative_prompt" class="text_pole textarea_compact" rows="3"></textarea>
+        </div>
+    </div>
+    <div class="inline-drawer">
+        <div class="inline-drawer-toggle inline-drawer-header">
+            <b>SD Prompt Templates</b>
+            <div class="inline-drawer-icon fa-solid fa-circle-chevron-down down"></div>
+        </div>
+        <div id="sd_prompt_templates" class="inline-drawer-content">
+        </div>
+    </div>
+</div>
diff --git a/server.js b/server.js
index fee94db2c..509d2adb4 100644
--- a/server.js
+++ b/server.js
@@ -4388,7 +4388,7 @@ app.post('/horde_generateimage', jsonParser, async (request, response) => {
         const ai_horde = getHordeClient();
         const generation = await ai_horde.postAsyncImageGenerate(
             {
-                prompt: `${request.body.prompt_prefix} ${request.body.prompt} ### ${request.body.negative_prompt}`,
+                prompt: `${request.body.prompt} ### ${request.body.negative_prompt}`,
                 params:
                 {
                     sampler_name: request.body.sampler,
@@ -4442,6 +4442,238 @@ app.post('/horde_generateimage', jsonParser, async (request, response) => {
     }
 });
 
+app.post('/api/novelai/generate-image', jsonParser, async (request, response) => {
+    if (!request.body) {
+        return response.sendStatus(400);
+    }
+
+    const key = readSecret(SECRET_KEYS.NOVEL);
+
+    if (!key) {
+        return response.sendStatus(401);
+    }
+
+    try {
+        console.log('NAI Diffusion request:', request.body);
+        const url = `${API_NOVELAI}/ai/generate-image`;
+        const result = await fetch(url, {
+            method: 'POST',
+            headers: {
+                'Authorization': `Bearer ${key}`,
+                'Content-Type': 'application/json',
+            },
+            body: JSON.stringify({
+                action: 'generate',
+                input: request.body.prompt,
+                model: request.body.model ?? 'nai-diffusion',
+                parameters: {
+                    negative_prompt: request.body.negative_prompt ?? '',
+                    height: request.body.height ?? 512,
+                    width: request.body.width ?? 512,
+                    scale: request.body.scale ?? 9,
+                    seed: Math.floor(Math.random() * 9999999999),
+                    sampler: request.body.sampler ?? 'k_dpmpp_2m',
+                    steps: request.body.steps ?? 28,
+                    n_samples: 1,
+                    // NAI handholding for prompts
+                    ucPreset: 0,
+                    qualityToggle: false,
+                },
+            }),
+        });
+
+        if (!result.ok) {
+            return response.sendStatus(500);
+        }
+
+        const archiveBuffer = await result.arrayBuffer();
+
+        const imageBuffer = await new Promise((resolve, reject) => yauzl.fromBuffer(Buffer.from(archiveBuffer), { lazyEntries: true }, (err, zipfile) => {
+            if (err) {
+                reject(err);
+            }
+
+            zipfile.readEntry();
+            zipfile.on('entry', (entry) => {
+                if (entry.fileName.endsWith('.png')) {
+                    console.log(`Extracting ${entry.fileName}`);
+                    zipfile.openReadStream(entry, (err, readStream) => {
+                        if (err) {
+                            reject(err);
+                        } else {
+                            const chunks = [];
+                            readStream.on('data', (chunk) => {
+                                chunks.push(chunk);
+                            });
+
+                            readStream.on('end', () => {
+                                const buffer = Buffer.concat(chunks);
+                                resolve(buffer);
+                                zipfile.readEntry(); // Continue to the next entry
+                            });
+                        }
+                    });
+                } else {
+                    zipfile.readEntry(); // Continue to the next entry
+                }
+            });
+        }));
+
+        const base64 = imageBuffer.toString('base64');
+        return response.send(base64);
+    } catch (error) {
+        console.log(error);
+        return response.sendStatus(500);
+    }
+});
+
+app.post('/api/sd/ping', jsonParser, async (request, response) => {
+    try {
+        const url = new URL(request.body.url);
+        url.pathname = '/internal/ping';
+
+        const result = await fetch(url);
+
+        if (!result.ok) {
+            throw new Error('SD WebUI returned an error.');
+        }
+
+        return response.sendStatus(200);
+    } catch (error) {
+        console.log(error);
+        return response.sendStatus(500);
+    }
+});
+
+app.post('/api/sd/samplers', jsonParser, async (request, response) => {
+    try {
+        const url = new URL(request.body.url);
+        url.pathname = '/sdapi/v1/samplers';
+
+        const result = await fetch(url);
+
+        if (!result.ok) {
+            throw new Error('SD WebUI returned an error.');
+        }
+
+        const data = await result.json();
+        const names = data.map(x => x.name);
+        return response.send(names);
+
+    } catch (error) {
+        console.log(error);
+        return response.sendStatus(500);
+    }
+});
+
+app.post('/api/sd/models', jsonParser, async (request, response) => {
+    try {
+        const url = new URL(request.body.url);
+        url.pathname = '/sdapi/v1/sd-models';
+
+        const result = await fetch(url);
+
+        if (!result.ok) {
+            throw new Error('SD WebUI returned an error.');
+        }
+
+        const data = await result.json();
+        const models = data.map(x => ({ value: x.title, text: x.title }));
+        return response.send(models);
+    } catch (error) {
+        console.log(error);
+        return response.sendStatus(500);
+    }
+});
+
+app.post('/api/sd/get-model', jsonParser, async (request, response) => {
+    try {
+        const url = new URL(request.body.url);
+        url.pathname = '/sdapi/v1/options';
+
+        const result = await fetch(url);
+        const data = await result.json();
+        return response.send(data['sd_model_checkpoint']);
+    } catch (error) {
+        console.log(error);
+        return response.sendStatus(500);
+    }
+});
+
+app.post('/api/sd/set-model', jsonParser, async (request, response) => {
+    try {
+        async function getProgress() {
+            const url = new URL(request.body.url);
+            url.pathname = '/sdapi/v1/progress';
+
+            const result = await fetch(url);
+            const data = await result.json();
+            return data;
+        }
+
+        const url = new URL(request.body.url);
+        url.pathname = '/sdapi/v1/options';
+
+        const options = {
+            sd_model_checkpoint: request.body.model,
+        };
+
+        const result = await fetch(url, {
+            method: 'POST',
+            body: JSON.stringify(options),
+            headers: { 'Content-Type': 'application/json' },
+        });
+
+        if (!result.ok) {
+            throw new Error('SD WebUI returned an error.');
+        }
+
+        const MAX_ATTEMPTS = 10;
+        const CHECK_INTERVAL = 2000;
+
+        for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt++) {
+            const progressState = await getProgress();
+
+            const progress = progressState["progress"]
+            const jobCount = progressState["state"]["job_count"];
+            if (progress == 0.0 && jobCount === 0) {
+                break;
+            }
+
+            console.log(`Waiting for SD WebUI to finish model loading... Progress: ${progress}; Job count: ${jobCount}`);
+            await delay(CHECK_INTERVAL);
+        }
+
+        return response.sendStatus(200);
+    } catch (error) {
+        console.log(error);
+        return response.sendStatus(500);
+    }
+});
+
+app.post('/api/sd/generate', jsonParser, async (request, response) => {
+    try {
+        const url = new URL(request.body.url);
+        url.pathname = '/sdapi/v1/txt2img';
+
+        const result = await fetch(url, {
+            method: 'POST',
+            body: JSON.stringify(request.body),
+            headers: { 'Content-Type': 'application/json' },
+        });
+
+        if (!result.ok) {
+            throw new Error('SD WebUI returned an error.');
+        }
+
+        const data = await result.json();
+        return response.send(data);
+    } catch (error) {
+        console.log(error);
+        return response.sendStatus(500);
+    }
+});
+
 app.post('/libre_translate', jsonParser, async (request, response) => {
     const key = readSecret(SECRET_KEYS.LIBRE);
     const url = readSecret(SECRET_KEYS.LIBRE_URL);