Merge pull request #1371 from LenAnderson/staging

Add ComfyUI to the Stable Diffusion extension

public/scripts/extensions/stable-diffusion/comfyWorkflowEditor.html

@@ -0,0 +1,30 @@
+<div id="sd_comfy_workflow_editor_template">
+	<div class="sd_comfy_workflow_editor">
+		<h3><strong>ComfyUI Workflow Editor</strong></h3>
+		<div class="sd_comfy_workflow_editor_content">
+			<div class="flex-container flexFlowColumn sd_comfy_workflow_editor_workflow_container">
+				<label for="sd_comfy_workflow_editor_workflow">Workflow (JSON)</label>
+				<textarea id="sd_comfy_workflow_editor_workflow" class="text_pole wide100p textarea_compact flex1" placeholder="Put the ComfyUI's workflow (JSON) here and replace the variable settings with placeholders."></textarea>
+			</div>
+			<div class="sd_comfy_workflow_editor_placeholder_container">
+				<div>Placeholders</div>
+				<ul class="sd_comfy_workflow_editor_placeholder_list">
+					<li data-placeholder="prompt" class="sd_comfy_workflow_editor_not_found">"%prompt%"</li>
+					<li data-placeholder="negative_prompt" class="sd_comfy_workflow_editor_not_found">"%negative_prompt%"</li>
+					<li data-placeholder="model" class="sd_comfy_workflow_editor_not_found">"%model%"</li>
+					<li data-placeholder="sampler" class="sd_comfy_workflow_editor_not_found">"%sampler%"</li>
+					<li data-placeholder="scheduler" class="sd_comfy_workflow_editor_not_found">"%scheduler%"</li>
+					<li data-placeholder="steps" class="sd_comfy_workflow_editor_not_found">"%steps%"</li>
+					<li data-placeholder="scale" class="sd_comfy_workflow_editor_not_found">"%scale%"</li>
+					<li data-placeholder="width" class="sd_comfy_workflow_editor_not_found">"%width%"</li>
+					<li data-placeholder="height" class="sd_comfy_workflow_editor_not_found">"%height%"</li>
+					<li><hr></li>
+					<li data-placeholder="seed" class="sd_comfy_workflow_editor_not_found">
+						"%seed%"
+						<a href="javascript:;" class="notes-link"><span class="note-link-span" title="Will generate a new random seed in SillyTavern that is then used in the ComfyUI workflow.">?</span></a>
+					</li>
+				</ul>
+			</div>
+		</div>
+	</div>
+</div>

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

@@ -43,6 +43,7 @@ const sources = {
     novel: 'novel',
     vlad: 'vlad',
     openai: 'openai',
+    comfy: 'comfy',
 }
 
 const generationMode = {
@@ -171,6 +172,9 @@ const defaultSettings = {
     steps_step: 1,
     steps: 20,
 
+    // Scheduler
+    scheduler: 'normal',
+
     // Image dimensions (Width & Height)
     dimension_min: 64,
     dimension_max: 2048,
@@ -235,6 +239,97 @@ const defaultSettings = {
 
     style: 'Default',
     styles: defaultStyles,
+
+    // ComyUI settings
+    comfy_url: 'http://127.0.0.1:8188',
+    comfy_workflow: `
+        {
+            "3": {
+                "class_type": "KSampler",
+                "inputs": {
+                    "cfg": "%scale%",
+                    "denoise": 1,
+                    "latent_image": [
+                        "5",
+                        0
+                    ],
+                    "model": [
+                        "4",
+                        0
+                    ],
+                    "negative": [
+                        "7",
+                        0
+                    ],
+                    "positive": [
+                        "6",
+                        0
+                    ],
+                    "sampler_name": "%sampler%",
+                    "scheduler": "%scheduler%",
+                    "seed": 8566257,
+                    "steps": "%steps%"
+                }
+            },
+            "4": {
+                "class_type": "CheckpointLoaderSimple",
+                "inputs": {
+                    "ckpt_name": "%model%"
+                }
+            },
+            "5": {
+                "class_type": "EmptyLatentImage",
+                "inputs": {
+                    "batch_size": 1,
+                    "height": "%height%",
+                    "width": "%width%"
+                }
+            },
+            "6": {
+                "class_type": "CLIPTextEncode",
+                "inputs": {
+                    "clip": [
+                        "4",
+                        1
+                    ],
+                    "text": "%prompt%"
+                }
+            },
+            "7": {
+                "class_type": "CLIPTextEncode",
+                "inputs": {
+                    "clip": [
+                        "4",
+                        1
+                    ],
+                    "text": "%negative_prompt%"
+                }
+            },
+            "8": {
+                "class_type": "VAEDecode",
+                "inputs": {
+                    "samples": [
+                        "3",
+                        0
+                    ],
+                    "vae": [
+                        "4",
+                        2
+                    ]
+                }
+            },
+            "9": {
+                "class_type": "SaveImage",
+                "inputs": {
+                    "filename_prefix": "SillyTavern",
+                    "images": [
+                        "8",
+                        0
+                    ]
+                }
+            }
+        }
+    `,
 }
 
 function processTriggers(chat, _, abort) {
@@ -371,6 +466,8 @@ async function loadSettings() {
     $('#sd_interactive_mode').prop('checked', extension_settings.sd.interactive_mode);
     $('#sd_openai_style').val(extension_settings.sd.openai_style);
     $('#sd_openai_quality').val(extension_settings.sd.openai_quality);
+    $('#sd_comfy_url').val(extension_settings.sd.comfy_url);
+    $('#sd_comfy_prompt').val(extension_settings.sd.comfy_prompt);
 
     for (const style of extension_settings.sd.styles) {
         const option = document.createElement('option');
@@ -383,7 +480,7 @@ async function loadSettings() {
     toggleSourceControls();
     addPromptTemplates();
 
-    await Promise.all([loadSamplers(), loadModels()]);
+    await Promise.all([loadSamplers(), loadModels(), loadSchedulers()]);
 }
 
 function addPromptTemplates() {
@@ -615,6 +712,11 @@ function onSamplerChange() {
     saveSettingsDebounced();
 }
 
+function onSchedulerChange() {
+    extension_settings.sd.scheduler = $('#sd_scheduler').find(':selected').val();
+    saveSettingsDebounced();
+}
+
 function onWidthInput() {
     extension_settings.sd.width = Number($('#sd_width').val());
     $('#sd_width_value').text(extension_settings.sd.width);
@@ -739,6 +841,17 @@ function onHrSecondPassStepsInput() {
     saveSettingsDebounced();
 }
 
+function onComfyUrlInput() {
+    extension_settings.sd.comfy_url = $('#sd_comfy_url').val();
+    saveSettingsDebounced();
+}
+
+function onComfyPromptInput() {
+    extension_settings.sd.comfy_prompt = $('#sd_comfy_prompt').val();
+    resetScrollHeight($(this));
+    saveSettingsDebounced();
+}
+
 async function validateAutoUrl() {
     try {
         if (!extension_settings.sd.auto_url) {
@@ -787,6 +900,33 @@ async function validateVladUrl() {
     }
 }
 
+async function validateComfyUrl() {
+    try {
+        if (!extension_settings.sd.comfy_url) {
+            throw new Error('URL is not set.');
+        }
+
+        
+        const result = await fetch(`/api/sd/comfy/ping`, {
+            method: 'POST',
+            headers: getRequestHeaders(),
+            body: JSON.stringify({
+                url: extension_settings.sd.comfy_url,
+            })
+        });
+        if (!result.ok) {
+            throw new Error('ComfyUI returned an error.');
+        }
+
+        await loadSamplers();
+        await loadSchedulers();
+        await loadModels();
+        toastr.success('ComfyUI API connected.');
+    } catch (error) {
+        toastr.error(`Could not validate ComfyUI API: ${error.message}`);
+    }
+}
+
 async function onModelChange() {
     extension_settings.sd.model = $('#sd_model').find(':selected').val();
     saveSettingsDebounced();
@@ -922,6 +1062,9 @@ async function loadSamplers() {
         case sources.openai:
             samplers = await loadOpenAiSamplers();
             break;
+        case sources.comfy:
+            samplers = await loadComfySamplers();
+            break;
     }
 
     for (const sampler of samplers) {
@@ -1031,6 +1174,29 @@ async function loadNovelSamplers() {
     ];
 }
 
+async function loadComfySamplers() {
+    if (!extension_settings.sd.comfy_url) {
+        return [];
+    }
+
+    try {
+        
+        const result = await fetch(`/api/sd/comfy/samplers`, {
+            method: 'POST',
+            headers: getRequestHeaders(),
+            body: JSON.stringify({
+                url: extension_settings.sd.comfy_url,
+            })
+        });
+        if (!result.ok) {
+            throw new Error('ComfyUI returned an error.');
+        }
+        return await result.json();
+    } catch (error) {
+        return [];
+    }
+}
+
 async function loadModels() {
     $('#sd_model').empty();
     let models = [];
@@ -1054,6 +1220,9 @@ async function loadModels() {
         case sources.openai:
             models = await loadOpenAiModels();
             break;
+        case sources.comfy:
+            models = await loadComfyModels();
+            break;
     }
 
     for (const model of models) {
@@ -1231,6 +1400,87 @@ async function loadNovelModels() {
     ];
 }
 
+async function loadComfyModels() {
+    if (!extension_settings.sd.comfy_url) {
+        return [];
+    }
+
+    try {
+        const result = await fetch(`/api/sd/comfy/models`, {
+            method: 'POST',
+            headers: getRequestHeaders(),
+            body: JSON.stringify({
+                url: extension_settings.sd.comfy_url,
+            })
+        });
+        if (!result.ok) {
+            throw new Error('ComfyUI returned an error.');
+        }
+        return await result.json();
+    } catch (error) {
+        return [];
+    }
+}
+
+async function loadSchedulers() {
+    $('#sd_scheduler').empty();
+    let schedulers = [];
+
+    switch (extension_settings.sd.source) {
+        case sources.extras:
+            schedulers = ['N/A'];
+            break;
+        case sources.horde:
+            schedulers = ['N/A'];
+            break;
+        case sources.auto:
+            schedulers = ['N/A'];
+            break;
+        case sources.novel:
+            schedulers = ['N/A'];
+            break;
+        case sources.vlad:
+            schedulers = ['N/A'];
+            break;
+        case sources.openai:
+            schedulers = ['N/A'];
+            break;
+        case sources.comfy:
+            schedulers = await loadComfySchedulers();
+            break;
+    }
+
+    for (const scheduler of schedulers) {
+        const option = document.createElement('option');
+        option.innerText = scheduler;
+        option.value = scheduler;
+        option.selected = scheduler === extension_settings.sd.scheduler;
+        $('#sd_scheduler').append(option);
+    }
+}
+
+async function loadComfySchedulers() {
+    if (!extension_settings.sd.comfy_url) {
+        return [];
+    }
+
+    try {
+        const result = await fetch(`/api/sd/comfy/schedulers`, {
+            method: 'POST',
+            headers: getRequestHeaders(),
+            body: JSON.stringify({
+                url: extension_settings.sd.comfy_url,
+            })
+        });
+        if (!result.ok) {
+            throw new Error('ComfyUI returned an error.');
+        }
+        return await result.json();
+    } catch (error) {
+        return [];
+    }
+}
+
 function getGenerationType(prompt) {
     let mode = generationMode.FREE;
 
@@ -1498,6 +1748,9 @@ async function sendGenerationRequest(generationType, prompt, characterName = nul
             case sources.openai:
                 result = await generateOpenAiImage(prefixedPrompt);
                 break;
+            case sources.comfy:
+                result = await generateComfyImage(prefixedPrompt);
+                break;
         }
 
         if (!result.data) {
@@ -1782,6 +2035,65 @@ async function generateOpenAiImage(prompt) {
     }
 }
 
+/**
+ * Generates an image in ComfyUI 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 generateComfyImage(prompt) {
+    const placeholders = [
+        'negative_prompt',
+        'model',
+        'sampler',
+        'scheduler',
+        'steps',
+        'scale',
+        'width',
+        'height',
+    ];
+
+    let workflow = extension_settings.sd.comfy_workflow.replace('"%prompt%"', JSON.stringify(prompt));
+    workflow = workflow.replace('"%seed%"', JSON.stringify(Math.round(Math.random()*Number.MAX_SAFE_INTEGER)));
+    placeholders.forEach(ph=>{
+        workflow = workflow.replace(`"%${ph}%"`, JSON.stringify(extension_settings.sd[ph]));
+    });
+    console.log(`{
+        "prompt": ${workflow}
+    }`);
+    const promptResult = await fetch(`/api/sd/comfy/generate`, {
+        method: 'POST',
+        headers: getRequestHeaders(),
+        body: JSON.stringify({
+            url: extension_settings.sd.comfy_url,
+            prompt: `{
+                "prompt": ${workflow}
+            }`,
+        })
+    });
+    return {format:'png', data:await promptResult.text()};
+}
+
+async function onComfyOpenWorkflowEditorClick() {
+    const editorHtml = $(await $.get('scripts/extensions/stable-diffusion/comfyWorkflowEditor.html'));
+    const popupResult = callPopup(editorHtml, "confirm", undefined, {okButton: "Save", wide:true, large:true, rows:1 });
+    const checkPlaceholders = ()=>{
+        const workflow = $('#sd_comfy_workflow_editor_workflow').val().toString();
+        $('.sd_comfy_workflow_editor_placeholder_list > li[data-placeholder]').each(function(idx) {
+            const key = this.getAttribute('data-placeholder');
+            const found = workflow.search(`"%${key}%"`) != -1;
+            this.classList[found?'remove':'add']('sd_comfy_workflow_editor_not_found');
+        });
+    };
+    $('#sd_comfy_workflow_editor_workflow').val(extension_settings.sd.comfy_workflow);
+    checkPlaceholders();
+    $('#sd_comfy_workflow_editor_workflow').on('input', checkPlaceholders);
+    if (await popupResult) {
+        extension_settings.sd.comfy_workflow = $('#sd_comfy_workflow_editor_workflow').val().toString();
+        saveSettingsDebounced();
+    }
+}
+
 async function sendMessage(prompt, image, generationType) {
     const context = getContext();
     const messageText = `[${context.name2} sends a picture that contains: ${prompt}]`;
@@ -1875,6 +2187,8 @@ function isValidState() {
             return secret_state[SECRET_KEYS.NOVEL];
         case sources.openai:
             return secret_state[SECRET_KEYS.OPENAI];
+        case sources.comfy:
+            return true;
     }
 }
 
@@ -1989,6 +2303,7 @@ jQuery(async () => {
     $('#sd_steps').on('input', onStepsInput);
     $('#sd_model').on('change', onModelChange);
     $('#sd_sampler').on('change', onSamplerChange);
+    $('#sd_scheduler').on('change', onSchedulerChange);
     $('#sd_prompt_prefix').on('input', onPromptPrefixInput);
     $('#sd_negative_prompt').on('input', onNegativePromptInput);
     $('#sd_width').on('input', onWidthInput);
@@ -2013,6 +2328,9 @@ jQuery(async () => {
     $('#sd_novel_upscale_ratio').on('input', onNovelUpscaleRatioInput);
     $('#sd_novel_anlas_guard').on('input', onNovelAnlasGuardInput);
     $('#sd_novel_view_anlas').on('click', onViewAnlasClick);
+    $('#sd_comfy_validate').on('click', validateComfyUrl);
+    $('#sd_comfy_url').on('input', onComfyUrlInput);
+    $('#sd_comfy_open_workflow_editor').on('click', onComfyOpenWorkflowEditorClick);
     $('#sd_expand').on('input', onExpandInput);
     $('#sd_style').on('change', onStyleSelect);
     $('#sd_save_style').on('click', onSaveStyleClick);

public/scripts/extensions/stable-diffusion/settings.html

@@ -34,6 +34,7 @@
                 <option value="vlad">SD.Next (vladmandic)</option>
                 <option value="novel">NovelAI Diffusion</option>
                 <option value="openai">OpenAI (DALL-E)</option>
+                <option value="comfy">ComfyUI</option>
             </select>
             <div data-sd-source="auto">
                 <label for="sd_auto_url">SD Web UI URL</label>
@@ -116,6 +117,23 @@
                     </select>
                 </div>
             </div>
+            <div data-sd-source="comfy">
+                <label for="sd_comfy_url">ComfyUI URL</label>
+                <div class="flex-container flexnowrap">
+                    <input id="sd_comfy_url" type="text" class="text_pole" placeholder="Example: {{comfy_url}}" value="{{comfy_url}}" />
+                    <div id="sd_comfy_validate" class="menu_button menu_button_icon">
+                        <i class="fa-solid fa-check"></i>
+                        <span data-i18n="Connect">
+                            Connect
+                        </span>
+                    </div>
+                </div>
+                <div id="sd_comfy_open_workflow_editor" class="menu_button">
+                    <i class="fa-solid fa-pen-to-square"></i>
+                    <span>Open Workflow Editor</span>
+                </div>
+                <p><i><b>Important:</b> The server must be accessible from the SillyTavern host machine.</i></p>
+            </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>
@@ -128,6 +146,10 @@
             <select id="sd_model"></select>
             <label for="sd_sampler">Sampling method</label>
             <select id="sd_sampler"></select>
+            <div data-sd-source="comfy">
+                <label for="sd_scheduler">Scheduler</label>
+                <select id="sd_scheduler"></select>
+            </div>
             <div class="flex-container marginTop10 margin-bot-10px">
                 <label class="flex1 checkbox_label">
                     <input id="sd_restore_faces" type="checkbox" />

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

@@ -27,3 +27,48 @@
     z-index: 30000;
     backdrop-filter: blur(--SmartThemeBlurStrength);
 }
+
+#sd_comfy_open_workflow_editor {
+    display: flex;
+    flex-direction: row;
+    gap: 10px;
+    width: fit-content;
+}
+#sd_comfy_workflow_editor_template {
+    height: 100%;
+}
+.sd_comfy_workflow_editor {
+    display: flex;
+    flex-direction: column;
+    height: 100%;
+}
+.sd_comfy_workflow_editor_content {
+    display: flex;
+    flex: 1 1 auto;
+    flex-direction: row;
+}
+.sd_comfy_workflow_editor_workflow_container {
+    flex: 1 1 auto;
+}
+#sd_comfy_workflow_editor_workflow {
+    font-family: monospace;
+}
+.sd_comfy_workflow_editor_placeholder_container {
+    flex: 0 0 auto;
+}
+.sd_comfy_workflow_editor_placeholder_list {
+    font-size: x-small;
+    list-style: none;
+    margin: 5px 0;
+    padding: 3px 5px;
+    text-align: left;
+}
+.sd_comfy_workflow_editor_placeholder_list > li[data-placeholder]:before {
+    content: "✅ ";
+}
+.sd_comfy_workflow_editor_placeholder_list > li.sd_comfy_workflow_editor_not_found:before {
+    content: "❌ ";
+}
+.sd_comfy_workflow_editor_placeholder_list > li > .notes-link {
+    cursor: help;
+}

src/stable-diffusion.js

@@ -347,6 +347,122 @@ function registerEndpoints(app, jsonParser) {
             return response.send({ prompt: originalPrompt });
         }
     });
+
+
+    app.post('/api/sd/comfy/ping', jsonParser, async(request, response)=>{
+		try {
+			const url = new URL(request.body.url);
+			url.pathname = '/system_stats'
+
+			const result = await fetch(url);
+			if (!result.ok) {
+				throw new Error('ComfyUI returned an error.');
+			}
+			
+			return response.sendStatus(200);
+		} catch (error) {
+			console.log(error);
+			return response.sendStatus(500);
+		}
+	});
+
+    app.post('/api/sd/comfy/samplers', jsonParser, async(request, response)=>{
+		try {
+			const url = new URL(request.body.url);
+			url.pathname = '/object_info'
+
+			const result = await fetch(url);
+			if (!result.ok) {
+				throw new Error('ComfyUI returned an error.');
+			}
+
+			const data = await result.json();
+			return response.send(data.KSampler.input.required.sampler_name[0]);
+		} catch (error) {
+			console.log(error);
+			return response.sendStatus(500);
+		}
+	});
+
+    app.post('/api/sd/comfy/models', jsonParser, async(request, response)=>{
+		try {
+			const url = new URL(request.body.url);
+			url.pathname = '/object_info'
+
+			const result = await fetch(url);
+			if (!result.ok) {
+				throw new Error('ComfyUI returned an error.');
+			}
+            const data = await result.json();
+			return response.send(data.CheckpointLoaderSimple.input.required.ckpt_name[0].map(it=>({value:it,text:it})));
+		} catch (error) {
+			console.log(error);
+			return response.sendStatus(500);
+		}
+	});
+
+    app.post('/api/sd/comfy/schedulers', jsonParser, async(request, response)=>{
+		try {
+			const url = new URL(request.body.url);
+			url.pathname = '/object_info'
+
+			const result = await fetch(url);
+			if (!result.ok) {
+				throw new Error('ComfyUI returned an error.');
+			}
+
+			const data = await result.json();
+			return response.send(data.KSampler.input.required.scheduler[0]);
+		} catch (error) {
+			console.log(error);
+			return response.sendStatus(500);
+		}
+	});
+
+    app.post('/api/sd/comfy/generate', jsonParser, async(request, response)=>{
+		try {
+			const url = new URL(request.body.url);
+			url.pathname = '/prompt'
+
+			const promptResult = await fetch(url, {
+                method: 'POST',
+                body: request.body.prompt,
+            });
+			if (!promptResult.ok) {
+				throw new Error('ComfyUI returned an error.');
+			}
+
+			const data = await promptResult.json();
+            const id = data.prompt_id;
+            let item;
+            const historyUrl = new URL(request.body.url);
+            historyUrl.pathname = '/history';
+            while (true) {
+                const result = await fetch(historyUrl);
+                if (!result.ok) {
+                    throw new Error('ComfyUI returned an error.');
+                }
+                const history = await result.json();
+                item = history[id];
+                if (item) {
+                    break;
+                }
+                await delay(100);
+            }
+            const imgInfo = Object.keys(item.outputs).map(it=>item.outputs[it].images).flat()[0];
+            const imgUrl = new URL(request.body.url);
+            imgUrl.pathname = '/view';
+            imgUrl.search = `?filename=${imgInfo.filename}&subfolder=${imgInfo.subfolder}&type=${imgInfo.type}`;
+            const imgResponse = await fetch(imgUrl);
+            if (!imgResponse.ok) {
+                throw new Error('ComfyUI returned an error.');
+            }
+            const imgBuffer = await imgResponse.buffer();
+            return response.send(imgBuffer.toString('base64'));
+		} catch (error) {
+			return response.sendStatus(500);
+		}
+	});
 }
 
 module.exports = {