Merge pull request #1371 from LenAnderson/staging

Add ComfyUI to the Stable Diffusion extension

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);