Comfyui serverless runpod image generation (#4891) * Add comfyrunpod generation endpoints. Get a bunch of settings added. Probably get runpod api secret stored properly. sending to runpod. Correct workflow wrapping for runpod. * Reuse comfyui workflow setting block for runpod. * Update API key url to point to runpod settings * Remove debug print. * Combine normal and runpod comfyui generation functions * Remove unused server endpoints for runpod. * Remove debug prints from server. * Remove redundant switch case. * Check for url to see if image generation is possible for comfy and comfy runpod. * lint * Correct/update doc comments. * Make runpod serverless a sub-type under comfyui source in the UI. * Set comfy_type in stable-diffusion defaultSettings * Allow loading comfy workflows even if url isnt set yet. * Change default runpod url to template example * Rename runpod api key to comfy_runpod * Handle default switches separately from standard now that default setting is ensured. * Update runpod api key name on server. * Move example runpod url to placeholder. * lint * Remove incorrect comment. * Hide toasts remote models toasts * React to key updates * Add type annotation for data and improve workflow handling in generate endpoint --------- Co-authored-by: Cohee <18619528+Cohee1207@users.noreply.github.com>

3098c4c5bff43640f4fb8baaaad0ae75d634053d

9nbf7c4q6b-lgtm <9nbf7c4q6b@privaterelay.appleid.com>

Signed

Parents: bff6ab3cb72a25473950b6d402147ac06ac8c7b4

public/scripts/extensions/stable-diffusion/index.js+138 -41 Viewed

     zai: 'zai',
     openrouter: 'openrouter',
 };
+const comfyTypes = {
+    standard: 'standard',
+    runpod_serverless: 'runpod_serverless',
+};
 const initiators = {
     command: 'command',

     styles: defaultStyles,
     // ComyUI settings
+    comfy_type: 'standard',
     comfy_url: 'http://127.0.0.1:8188',
     comfy_workflow: 'Default_Comfy_Workflow.json',
+    comfy_runpod_url: '',
     // Pollinations settings
     pollinations_enhance: false,

         const source = $(this).data('sd-source').split(',');
         $(this).toggle(source.includes(extension_settings.sd.source));
     });
+    $('.sd_settings [data-sd-comfy-type]').each(function () {
+        const source = $(this).data('sd-comfy-type').split(',');
+        $(this).toggle(source.includes(extension_settings.sd.comfy_type));
+    });
+}
 async function loadSettings() {

     $('#sd_openai_style').val(extension_settings.sd.openai_style);
     $('#sd_openai_quality').val(extension_settings.sd.openai_quality);
     $('#sd_openai_duration').val(extension_settings.sd.openai_duration);
+    $('#sd_comfy_type').val(extension_settings.sd.comfy_type);
     $('#sd_comfy_url').val(extension_settings.sd.comfy_url);
     $('#sd_comfy_prompt').val(extension_settings.sd.comfy_prompt);
+    $('#sd_comfy_runpod_url').val(extension_settings.sd.comfy_runpod_url);
     $('#sd_snap').prop('checked', extension_settings.sd.snap);
     $('#sd_clip_skip').val(extension_settings.sd.clip_skip);
     $('#sd_clip_skip_value').val(extension_settings.sd.clip_skip);

     await loadSettingOptions();
+}
+async function onComfyTypeChange() {
+    extension_settings.sd.comfy_type = $('#sd_comfy_type').find(':selected').val();
+    await onSourceChange();
+}
 function onFunctionToolInput() {
     extension_settings.sd.function_tool = !!$(this).prop('checked');
     saveSettingsDebounced();

+}
 function onComfyUrlInput() {
-    // Remove trailing slashes
     extension_settings.sd.comfy_url = String($('#sd_comfy_url').val());
     saveSettingsDebounced();
+}
+function onComfyRunPodUrlInput() {
+    extension_settings.sd.comfy_runpod_url = String($('#sd_comfy_runpod_url').val());
+    saveSettingsDebounced();
+}
 function onHFModelInput() {
     extension_settings.sd.huggingface_model_id = $('#sd_huggingface_model_id').val();
     saveSettingsDebounced();

+    }
+}
+async function validateComfyRunPodUrl() {
+    try {
+        if (!extension_settings.sd.comfy_runpod_url) {
+            throw new Error('URL is not set.');
+        }
+        const result = await fetch('/api/sd/comfyrunpod/ping', {
+            method: 'POST',
+            headers: getRequestHeaders(),
+            body: JSON.stringify({
+                url: extension_settings.sd.comfy_runpod_url,
+            }),
+        });
+        if (!result.ok) {
+            throw new Error('ComfyUI RunPod returned an error.');
+        }
+        await loadSettingOptions();
+        toastr.success('ComfyUI RunPod API connected.');
+    } catch (error) {
+        toastr.error(`Could not validate ComfyUI RunPod API: ${error.message}`);
+    }
+}
 async function onModelChange() {
     const selectedModel = $('#sd_model').find(':selected');
     extension_settings.sd.model = selectedModel.val();

     switchModelSpecificControls(extension_settings.sd.model);
-    const cloudSources = [
+    const updateRemoteModelSources = [
-        sources.horde,
+        sources.auto,
-        sources.novel,
+        sources.vlad,
-        sources.openai,
+        sources.extras,
-        sources.aimlapi,
-        sources.togetherai,
-        sources.pollinations,
-        sources.stability,
-        sources.huggingface,
-        sources.electronhub,
-        sources.nanogpt,
-        sources.bfl,
-        sources.falai,
-        sources.xai,
-        sources.google,
-        sources.chutes,
-        sources.zai,
-        sources.openrouter,
     ];
-    if (cloudSources.includes(extension_settings.sd.source)) {
+    if (!updateRemoteModelSources.includes(extension_settings.sd.source)) {
         return;
+    }

@@ -1685,6 +1718,9 @@ async function loadNovelSamplers() {
1685	}	1718	}
1686		1719
1687	async function loadComfySamplers() {	1720	async function loadComfySamplers() {
		1721	if (extension_settings.sd.comfy_type === comfyTypes.runpod_serverless) {
		1722	return ['N/A'];
		1723	}
1688	if (!extension_settings.sd.comfy_url) {	1724	if (!extension_settings.sd.comfy_url) {
1689	return [];	1725	return [];
1690	}	1726	}

+}
 async function loadComfyModels() {
+    if (extension_settings.sd.comfy_type === comfyTypes.runpod_serverless) {
+        $('#sd_runpod_key').toggleClass('success', !!secret_state[SECRET_KEYS.COMFY_RUNPOD]);
+        return [
+            { value: '', text: 'N/A' },
+        ];
+    }
     if (!extension_settings.sd.comfy_url) {
         return [];
+    }

@@ -2388,6 +2430,9 @@ async function loadSchedulers() {
2388	}	2430	}
2389		2431
2390	async function loadComfySchedulers() {	2432	async function loadComfySchedulers() {
		2433	if (extension_settings.sd.comfy_type === comfyTypes.runpod_serverless) {
		2434	return ['N/A'];
		2435	}
2391	if (!extension_settings.sd.comfy_url) {	2436	if (!extension_settings.sd.comfy_url) {
2392	return [];	2437	return [];
2393	}	2438	}

@@ -2521,6 +2566,9 @@ async function loadAutoVaes() {
2521	}	2566	}
2522		2567
2523	async function loadComfyVaes() {	2568	async function loadComfyVaes() {
		2569	if (extension_settings.sd.comfy_type === comfyTypes.runpod_serverless) {
		2570	return ['N/A'];
		2571	}
2524	if (!extension_settings.sd.comfy_url) {	2572	if (!extension_settings.sd.comfy_url) {
2525	return [];	2573	return [];
2526	}	2574	}

+}
 async function loadComfyWorkflows() {
-    if (!extension_settings.sd.comfy_url) {
-        return;
-    }
     try {
         $('#sd_comfy_workflow').empty();
         const result = await fetch('/api/sd/comfy/workflows', {

                 result = await generateAimlapiImage(prefixedPrompt, signal);
                 break;
             case sources.comfy:
-                result = await generateComfyImage(prefixedPrompt, negativePrompt, signal);
+                switch (extension_settings.sd.comfy_type) {
+                    case comfyTypes.runpod_serverless:
+                        result = await generateComfyRunPodImage(prefixedPrompt, negativePrompt, signal);
+                        break;
+                    case comfyTypes.standard:
+                        result = await generateComfyImage(prefixedPrompt, negativePrompt, signal);
+                        break;
+                    default:
+                        throw new Error('Unknown comfyUI server type.');
+                }
                 break;
             case sources.togetherai:
                 result = await generateTogetherAIImage(prefixedPrompt, negativePrompt, signal);

+}
 /**
- * Generates an image in ComfyUI using the provided prompt and configuration settings.
+ * Generates an image in local ComfyUI or serverless runpod using the provided prompt and configuration settings.
+ *
  * @param {string} prompt - The main instruction used to guide the image generation.
  * @param {string} negativePrompt - The instruction used to restrict the image generation.
  * @param {AbortSignal} signal - An AbortSignal object that can be used to cancel the request.
+ * @param {string} basePath - ST server endpoint for the service. '/api/sd/comfy' for local, '/api/sd/comfyrunpod' for serverless.
+ * @param {string[]} placeholders - Array of substitutions to apply to the workflow.
+ * @param {string} url - The url of the service to call. Passed to ST server.
  * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete.
  */
-async function generateComfyImage(prompt, negativePrompt, signal) {
+async function generateComfyImageCommon(prompt, negativePrompt, signal, basePath, placeholders, url) {
-    const placeholders = [
-        'model',
-        'vae',
-        'sampler',
-        'scheduler',
-        'steps',
-        'scale',
-        'width',
-        'height',
-    ];
     const workflowResponse = await fetch('/api/sd/comfy/workflow', {
         method: 'POST',
         headers: getRequestHeaders(),

     console.log(`{
         "prompt": ${workflow}
     }`);
-    const promptResult = await fetch('/api/sd/comfy/generate', {
+    const promptResult = await fetch(`${basePath}/generate`, {
         method: 'POST',
         headers: getRequestHeaders(),
         signal: signal,
         body: JSON.stringify({
-            url: extension_settings.sd.comfy_url,
+            url,
             prompt: `{
                 "prompt": ${workflow}
             }`,

     return { format, data };
+}
+/**
+ * Generates an image in ComfyUI using the provided prompt and configuration settings.
+ *
+ * @param {string} prompt - The main instruction used to guide the image generation.
+ * @param {string} negativePrompt - The instruction used to restrict the image generation.
+ * @param {AbortSignal} signal - An AbortSignal object that can be used to cancel the request.
+ * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete.
+ */
+async function generateComfyImage(prompt, negativePrompt, signal) {
+    const placeholders = [
+        'model',
+        'vae',
+        'sampler',
+        'scheduler',
+        'steps',
+        'scale',
+        'width',
+        'height',
+    ];
+    return generateComfyImageCommon(prompt, negativePrompt, signal, '/api/sd/comfy', placeholders, extension_settings.sd.comfy_url);
+}
+/**
+ * Generates an image using ComfyUI through serverless runpod endpoint using the provided prompt and configuration settings.
+ *
+ * @param {string} prompt - The main instruction used to guide the image generation.
+ * @param {string} negativePrompt - The instruction used to restrict the image generation.
+ * @param {AbortSignal} signal - An AbortSignal object that can be used to cancel the request.
+ * @returns {Promise<{format: string, data: string}>} - A promise that resolves when the image generation and processing are complete.
+ */
+async function generateComfyRunPodImage(prompt, negativePrompt, signal) {
+    const placeholders = [
+        'steps',
+        'scale',
+        'width',
+        'height',
+    ];
+    return generateComfyImageCommon(prompt, negativePrompt, signal, '/api/sd/comfyrunpod', placeholders, extension_settings.sd.comfy_runpod_url);
+}
 /**
  * Generates an image in Hugging Face Inference API using the provided prompt and configuration settings (model selected).

         case sources.aimlapi:
             return secret_state[SECRET_KEYS.AIMLAPI];
         case sources.comfy:
-            return true;
+            switch (extension_settings.sd.comfy_type) {
+                case comfyTypes.runpod_serverless:
+                    return !!extension_settings.sd.comfy_runpod_url &&
+                        secret_state[SECRET_KEYS.COMFY_RUNPOD];
+                case comfyTypes.standard:
+                    return !!extension_settings.sd.comfy_url;
+                default:
+                    return false;
+            }
         case sources.togetherai:
             return secret_state[SECRET_KEYS.TOGETHERAI];
         case sources.pollinations:

     $('#sd_novel_decrisper').on('input', onNovelDecrisperInput);
     $('#sd_novel_variety_boost').on('input', onNovelVarietyBoostInput);
     $('#sd_pollinations_enhance').on('input', onPollinationsEnhanceInput);
+    $('#sd_comfy_type').on('change', onComfyTypeChange);
     $('#sd_comfy_validate').on('click', validateComfyUrl);
+    $('#sd_comfy_runpod_validate').on('click', validateComfyRunPodUrl);
     $('#sd_comfy_url').on('input', onComfyUrlInput);
+    $('#sd_comfy_runpod_url').on('input', onComfyRunPodUrlInput);
     $('#sd_comfy_workflow').on('change', onComfyWorkflowChange);
     $('#sd_comfy_open_workflow_editor').on('click', onComfyOpenWorkflowEditorClick);
     $('#sd_comfy_new_workflow').on('click', onComfyNewWorkflowClick);

                 case SECRET_KEYS.FALAI:
                 case SECRET_KEYS.STABILITY:
                 case SECRET_KEYS.AIMLAPI:
+                case SECRET_KEYS.COMFY_RUNPOD:
                     await loadSettingOptions();
                     break;
+            }

public/scripts/extensions/stable-diffusion/settings.html+38 -8 Viewed

                 </div>
             </div>
             <div data-sd-source="comfy">
-                <label for="sd_comfy_url">ComfyUI URL</label>
+                <label for="sd_comfy_type">Server Type</label>
-                <div class="flex-container flexnowrap">
+                <select id="sd_comfy_type">
-                    <input id="sd_comfy_url" type="text" class="text_pole" data-i18n="[placeholder]sd_comfy_url" placeholder="Example: {{comfy_url}}" value="{{comfy_url}}" />
+                    <option value="standard">Standard Server</option>
-                    <div id="sd_comfy_validate" class="menu_button menu_button_icon">
+                    <option value="runpod_serverless">RunPod Serverless Endpoint</option>
-                        <i class="fa-solid fa-check"></i>
+                </select>
-                        <span data-i18n="Connect">
+                <div data-sd-comfy-type="standard">
-                            Connect
+                    <label for="sd_comfy_url">ComfyUI URL</label>
-                        </span>
+                    <div class="flex-container flexnowrap">
+                        <input id="sd_comfy_url" type="text" class="text_pole" data-i18n="[placeholder]sd_comfy_url" 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>
+                <div data-sd-comfy-type="runpod_serverless">
+                    <label for="sd_comfy_runpod_url">ComfyUI RunPod URL</label>
+                    <div class="flex-container flexnowrap">
+                        <input id="sd_comfy_runpod_url" type="text" class="text_pole" data-i18n="[placeholder]sd_comfy_runpod_url" placeholder="eg: https://api.runpod.ai/v2/<your endpoint id>" value="{{comfy_runpod_url}}" />
+                        <div id="sd_comfy_runpod_validate" class="menu_button menu_button_icon">
+                            <i class="fa-solid fa-check"></i>
+                            <span data-i18n="Connect">
+                                Connect
+                            </span>
+                        </div>
+                    </div>
+                    <div class="flex-container flexnowrap alignItemsBaseline marginBot5">
+                        <a href="https://console.runpod.io/user/settings" target="_blank" rel="noopener noreferrer">
+                            <strong data-i18n="API Key">API Key</strong>
+                            <i class="fa-solid fa-share-from-square"></i>
+                        </a>
+                        <span class="expander"></span>
+                        <div id="sd_runpod_key" class="menu_button menu_button_icon manage-api-keys" data-key="api_key_comfy_runpod">
+                            <i class="fa-fw fa-solid fa-key"></i>
+                            <span data-i18n="Click to set">Click to set</span>
+                        </div>
                     </div>
                 </div>
                 <p><i><b data-i18n="Important:">Important:</b></i><i data-i18n="The server must be accessible from the SillyTavern host machine."> The server must be accessible from the SillyTavern host machine.</i></p>

public/scripts/secrets.js+3 -0 Viewed

     NANOGPT: 'api_key_nanogpt',
     TAVILY: 'api_key_tavily',
     BFL: 'api_key_bfl',
+    COMFY_RUNPOD: 'api_key_comfy_runpod',
     GENERIC: 'api_key_generic',
     DEEPSEEK: 'api_key_deepseek',
     SERPER: 'api_key_serper',

     [SECRET_KEYS.CUSTOM_OPENAI_TTS]: 'Custom OpenAI TTS',
     [SECRET_KEYS.TAVILY]: 'Tavily',
     [SECRET_KEYS.BFL]: 'Black Forest Labs',
+    [SECRET_KEYS.COMFY_RUNPOD]: 'ComfyUI RunPod',
     [SECRET_KEYS.SERPAPI]: 'SerpApi',
     [SECRET_KEYS.SERPER]: 'Serper',
     [SECRET_KEYS.FALAI]: 'FAL.AI',

     [SECRET_KEYS.AZURE_OPENAI]: '#api_key_azure_openai',
     [SECRET_KEYS.ZAI]: '#api_key_zai',
     [SECRET_KEYS.SILICONFLOW]: '#api_key_siliconflow',
+    [SECRET_KEYS.COMFY_RUNPOD]: '#api_key_comfy_runpod',
 };
 const getLabel = () => moment().format('L LT');

src/endpoints/secrets.js+1 -0 Viewed

     ELECTRONHUB: 'api_key_electronhub',
     NANOGPT: 'api_key_nanogpt',
     BFL: 'api_key_bfl',
+    COMFY_RUNPOD: 'api_key_comfy_runpod',
     FALAI: 'api_key_falai',
     GENERIC: 'api_key_generic',
     DEEPSEEK: 'api_key_deepseek',

src/endpoints/stable-diffusion.js+103 -0 Viewed

+    }
 });
+const comfyRunPod = express.Router();
+comfyRunPod.post('/ping', async (request, response) => {
+    try {
+        const key = readSecret(request.user.directories, SECRET_KEYS.COMFY_RUNPOD);
+        if (!key) {
+            console.warn('RunPod key not found.');
+            return response.sendStatus(400);
+        }
+        const url = new URL(urlJoin(request.body.url, '/health'));
+        const result = await fetch(url, {
+            method: 'GET',
+            headers: { 'Authorization': `Bearer ${key}` },
+        });
+        if (!result.ok) {
+            throw new Error('ComfyUI returned an error.');
+        }
+        /** @type {any} */
+        const data = await result.json();
+        if (data.workers.ready <= 0) {
+            console.warn(`No workers reported as ready. ${result}`);
+        }
+        return response.sendStatus(200);
+    } catch (error) {
+        console.error(error);
+        return response.sendStatus(500);
+    }
+});
+comfyRunPod.post('/generate', async (request, response) => {
+    try {
+        const key = readSecret(request.user.directories, SECRET_KEYS.COMFY_RUNPOD);
+        if (!key) {
+            console.warn('RunPod key not found.');
+            return response.sendStatus(400);
+        }
+        let jobId;
+        let item;
+        const url = new URL(urlJoin(request.body.url, '/run'));
+        const controller = new AbortController();
+        request.socket.removeAllListeners('close');
+        request.socket.on('close', function () {
+            if (!response.writableEnded && !item) {
+                const interruptUrl = new URL(urlJoin(request.body.url, `/cancel/${jobId}`));
+                fetch(interruptUrl, { method: 'POST', headers: { 'Authorization': `Bearer ${key}` } });
+            }
+            controller.abort();
+        });
+        const workflow = JSON.parse(request.body.prompt).prompt;
+        const wrappedWorkflow = workflow?.input?.workflow ? workflow : ({ input: { workflow: workflow } });
+        const runpodPrompt = JSON.stringify(wrappedWorkflow);
+        console.debug('ComfyUI RunPod request:', wrappedWorkflow);
+        const promptResult = await fetch(url, {
+            method: 'POST',
+            headers: { 'Authorization': `Bearer ${key}` },
+            body: runpodPrompt,
+        });
+        if (!promptResult.ok) {
+            const text = await promptResult.text();
+            throw new Error('ComfyUI returned an error.', { cause: tryParse(text) });
+        }
+        /** @type {any} */
+        const data = await promptResult.json();
+        jobId = data.id;
+        const statusUrl = new URL(urlJoin(request.body.url, `/status/${jobId}`));
+        while (true) {
+            const result = await fetch(statusUrl, {
+                method: 'GET',
+                headers: { 'Authorization': `Bearer ${key}` },
+            });
+            if (!result.ok) {
+                throw new Error('ComfyUI returned an error.');
+            }
+            /** @type {any} */
+            const status = await result.json();
+            if (status.output) {
+                item = status.output.images[0];
+            }
+            if (item) {
+                break;
+            }
+            await delay(500);
+        }
+        const format = path.extname(item.filename).slice(1).toLowerCase() || 'png';
+        return response.send({ format: format, data: item.data });
+    } catch (error) {
+        console.error('ComfyUI error:', error);
+        response.status(500).send(error.message);
+        return response;
+    }
+});
 const together = express.Router();
 together.post('/models', async (request, response) => {

 });
 router.use('/comfy', comfy);
+router.use('/comfyrunpod', comfyRunPod);
 router.use('/together', together);
 router.use('/drawthings', drawthings);
 router.use('/pollinations', pollinations);