Image Generation: Support video outputs for ComfyUI (#4194) * ImgGen: ComfyUI support video outputs * Fix video format check to handle undefined and case sensitivity * Fix error handling and add output debug log in comfy workflows * Refactor Comfy generation result * Add video file extensions constant

7755fc50a5c6897b91d9edf1c5c4efe21a380806

Cohee <18619528+Cohee1207@users.noreply.github.com>

Signed
Parents: 9e6fbde3cfb064098dc61cd3cb9f6b3d97944ad4
4 files changed, +47 -18Ignore whitespace
public/scripts/constants.js+6 -0
@@ -22,3 +22,9 @@ export const debounce_timeout = {
2222 * which is needed to preserve world info timed effects.
2323 */
2424export const IGNORE_SYMBOL = Symbol.for('ignore');
25+
26+/**
27+ * Common video file extensions. Should be the same as supported by Gemini.
28+ * https://ai.google.dev/gemini-api/docs/video-understanding#supported-formats
29+ */
30+export const VIDEO_EXTENSIONS = ['mp4', 'avi', 'mov', 'wmv', 'flv', 'webm', '3gp', 'mkv'];
public/scripts/data-maid.js+2 -1
@@ -1,4 +1,5 @@
11import { getRequestHeaders } from '../script.js';
2+import { VIDEO_EXTENSIONS } from './constants.js';
23import { t } from './i18n.js';
34import { callGenericPopup, Popup, POPUP_TYPE } from './popup.js';
45import { renderTemplateAsync } from './templates.js';
@@ -350,7 +351,7 @@ class DataMaidDialog {
350351 * @private
351352 */
352353 async getViewElement(url, name) {
353- const isVideo = /\.(mp4|webm|ogg|avi|mov|3gp|flv|mkv|wmv)$/i.test(name);
354+ const isVideo = VIDEO_EXTENSIONS.includes(name.split('.').pop());
354355 const mediaElement = document.createElement(isVideo ? 'video' : 'img');
355356 if (mediaElement instanceof HTMLVideoElement) {
356357 mediaElement.controls = true;
public/scripts/extensions/stable-diffusion/index.js+30 -14
@@ -51,7 +51,7 @@ import {
5151 SlashCommandArgument,
5252 SlashCommandNamedArgument,
5353} from '../../slash-commands/SlashCommandArgument.js';
5454import { debounce_timeout, VIDEO_EXTENSIONS } from '../../constants.js';
5555import { SlashCommandEnumValue } from '../../slash-commands/SlashCommandEnumValue.js';
5656import { callGenericPopup, Popup, POPUP_TYPE } from '../../popup.js';
5757import { commonEnumProviders } from '../../slash-commands/SlashCommandCommonEnumsProvider.js';
@@ -339,6 +339,7 @@ const defaultSettings = {
339339};
340340
341341const writePromptFieldsDebounced = debounce(writePromptFields, debounce_timeout.relaxed);
342+const isVideo = (/** @type {string} */ format) => VIDEO_EXTENSIONS.includes(String(format || '').trim().toLowerCase());
342343
343344/**
344345 * Generate interceptor for interactive mode triggers.
@@ -2476,14 +2477,14 @@ async function generatePicture(initiator, args, trigger, message, callback) {
24762477
24772478 if (generationType === generationMode.BACKGROUND) {
24782479 const callbackOriginal = callback;
24792480 callback = async function (prompt, imagePath, generationType, _negativePromptPrefix, _initiator, prefixedPrompt, format) {
24802481 const imgUrl = `url("${encodeURI(imagePath)}")`;
24812482 await eventSource.emit(event_types.FORCE_SET_BACKGROUND, { url: imgUrl, path: imagePath });
24822483
24832484 if (typeof callbackOriginal === 'function') {
24842485 await callbackOriginal(prompt, imagePath, generationType, negativePromptPrefix, initiator, prefixedPrompt, format);
24852486 } else {
24862487 await sendMessage(prompt, imagePath, generationType, negativePromptPrefix, initiator, prefixedPrompt, format);
24872488 }
24882489 };
24892490 }
@@ -2838,8 +2839,8 @@ async function sendGenerationRequest(generationType, prompt, additionalNegativeP
28382839 const filename = `${characterName}_${humanizedDateTime()}`;
28392840 const base64Image = await saveBase64AsFile(result.data, characterName, filename, result.format);
28402841 callback
28412842 ? await callback(prompt, base64Image, generationType, additionalNegativePrefix, initiator, prefixedPrompt, result.format)
28422843 : await sendMessage(prompt, base64Image, generationType, additionalNegativePrefix, initiator, prefixedPrompt, result.format);
28432844 return base64Image;
28442845}
28452846
@@ -3524,7 +3525,8 @@ async function generateComfyImage(prompt, negativePrompt, signal) {
35243525 const text = await promptResult.text();
35253526 throw new Error(text);
35263527 }
35273528 returnconst { format: 'png', data: } = await promptResult.textjson() };
3529+ return { format, data };
35283530}
35293531
35303532
@@ -3864,8 +3866,9 @@ async function onComfyDeleteWorkflowClick() {
38643866 * @param {string} additionalNegativePrefix Additional negative prompt used for the image generation
38653867 * @param {string} initiator The initiator of the image generation
38663868 * @param {string} prefixedPrompt Prompt with an attached specific prefix
3869+ * @param {string} format Format of the image (e.g., 'png', 'jpg')
38673870 */
38683871async function sendMessage(prompt, image, generationType, additionalNegativePrefix, initiator, prefixedPrompt, format) {
38693872 const context = getContext();
38703873 const name = context.groupId ? systemUserName : context.name2;
38713874 const template = extension_settings.sd.prompts[generationMode.MESSAGE] || '{{prompt}}';
@@ -3885,6 +3888,12 @@ async function sendMessage(prompt, image, generationType, additionalNegativePref
38853888 image_swipes: [image],
38863889 },
38873890 };
3891+ if (isVideo(format)) {
3892+ message.extra.video = image;
3893+ delete message.extra.image;
3894+ delete message.extra.image_swipes;
3895+ delete message.extra.inline_image;
3896+ }
38883897 context.chat.push(message);
38893898 const messageId = context.chat.length - 1;
38903899 await eventSource.emit(event_types.MESSAGE_RECEIVED, messageId, 'extension');
@@ -4068,7 +4077,7 @@ async function sdMessageButton(e) {
40684077 }
40694078 }
40704079
40714080 function saveGeneratedImage(prompt, image, generationType, negative, _initiator, _prefixedPrompt, format) {
40724081 // Some message sources may not create the extra object
40734082 if (typeof message.extra !== 'object' || message.extra === null) {
40744083 message.extra = {};
@@ -4085,17 +4094,24 @@ async function sdMessageButton(e) {
40854094 swipes.push(message.extra.image);
40864095 }
40874096
4088- swipes.push(image);
4097+ const isVideoFormat = isVideo(format);
4098+
4099+ if (isVideoFormat) {
4100+ message.extra.video = image;
4101+ } else {
4102+ swipes.push(image);
4103+
4104+ // If already contains an image and it's not inline - leave it as is
4105+ message.extra.inline_image = !(message.extra.image && !message.extra.inline_image);
4106+ message.extra.image = image;
4107+ }
40894108
4090- // If already contains an image and it's not inline - leave it as is
4091- message.extra.inline_image = !(message.extra.image && !message.extra.inline_image);
4092- message.extra.image = image;
40934109 message.extra.title = prompt;
40944110 message.extra.generationType = generationType;
40954111 message.extra.negative = negative;
40964112 appendMediaToMessage(message, $mes);
40974113
40984114 return context.saveChat();
40994115 }
41004116}
41014117
src/endpoints/stable-diffusion.js+9 -3
@@ -440,7 +440,7 @@ comfy.post('/models', async (request, response) => {
440440 models.forEach(it => it.text = it.text.replace(/\.[^.]*$/, '').replace(/_/g, ' '));
441441
442442 return response.send(models);
443443 } catch (error) {
444444 console.error(error);
445445 return response.sendStatus(500);
446446 }
@@ -581,15 +581,21 @@ comfy.post('/generate', async (request, response) => {
581581 .join('\n') || '';
582582 throw new Error(`ComfyUI generation did not succeed.\n\n${errorMessages}`.trim());
583583 }
584584 const imgInfooutputs = Object.keys(item.outputs).map(it => item.outputs[it].images).flat()[0];
585+ console.debug('ComfyUI outputs:', outputs);
586+ const imgInfo = outputs.map(it => it.images).flat()[0] ?? outputs.map(it => it.gifs).flat()[0];
587+ if (!imgInfo) {
588+ throw new Error('ComfyUI did not return any recognizable outputs.');
589+ }
585590 const imgUrl = new URL(urlJoin(request.body.url, '/view'));
586591 imgUrl.search = `?filename=${imgInfo.filename}&subfolder=${imgInfo.subfolder}&type=${imgInfo.type}`;
587592 const imgResponse = await fetch(imgUrl);
588593 if (!imgResponse.ok) {
589594 throw new Error('ComfyUI returned an error.');
590595 }
596+ const format = path.extname(imgInfo.filename).slice(1).toLowerCase() || 'png';
591597 const imgBuffer = await imgResponse.arrayBuffer();
592598 return response.send({ format: format, data: Buffer.from(imgBuffer).toString('base64') });
593599 } catch (error) {
594600 console.error('ComfyUI error:', error);
595601 response.status(500).send(error.message);