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

public/scripts/constants.js+6 -0 Viewed

  * which is needed to preserve world info timed effects.
  */
 export const IGNORE_SYMBOL = Symbol.for('ignore');
+/**
+ * Common video file extensions. Should be the same as supported by Gemini.
+ * https://ai.google.dev/gemini-api/docs/video-understanding#supported-formats
+ */
+export const VIDEO_EXTENSIONS = ['mp4', 'avi', 'mov', 'wmv', 'flv', 'webm', '3gp', 'mkv'];

public/scripts/data-maid.js+2 -1 Viewed

 import { getRequestHeaders } from '../script.js';
+import { VIDEO_EXTENSIONS } from './constants.js';
 import { t } from './i18n.js';
 import { callGenericPopup, Popup, POPUP_TYPE } from './popup.js';
 import { renderTemplateAsync } from './templates.js';

      * @private
      */
     async getViewElement(url, name) {
-        const isVideo = /\.(mp4|webm|ogg|avi|mov|3gp|flv|mkv|wmv)$/i.test(name);
+        const isVideo = VIDEO_EXTENSIONS.includes(name.split('.').pop());
         const mediaElement = document.createElement(isVideo ? 'video' : 'img');
         if (mediaElement instanceof HTMLVideoElement) {
             mediaElement.controls = true;

public/scripts/extensions/stable-diffusion/index.js+30 -14 Viewed

     SlashCommandArgument,
     SlashCommandNamedArgument,
 } from '../../slash-commands/SlashCommandArgument.js';
-import { debounce_timeout } from '../../constants.js';
+import { debounce_timeout, VIDEO_EXTENSIONS } from '../../constants.js';
 import { SlashCommandEnumValue } from '../../slash-commands/SlashCommandEnumValue.js';
 import { callGenericPopup, Popup, POPUP_TYPE } from '../../popup.js';
 import { commonEnumProviders } from '../../slash-commands/SlashCommandCommonEnumsProvider.js';

 };
 const writePromptFieldsDebounced = debounce(writePromptFields, debounce_timeout.relaxed);
+const isVideo = (/** @type {string} */ format) => VIDEO_EXTENSIONS.includes(String(format || '').trim().toLowerCase());
 /**
  * Generate interceptor for interactive mode triggers.

     if (generationType === generationMode.BACKGROUND) {
         const callbackOriginal = callback;
-        callback = async function (prompt, imagePath, generationType, _negativePromptPrefix, _initiator, prefixedPrompt) {
+        callback = async function (prompt, imagePath, generationType, _negativePromptPrefix, _initiator, prefixedPrompt, format) {
             const imgUrl = `url("${encodeURI(imagePath)}")`;
             await eventSource.emit(event_types.FORCE_SET_BACKGROUND, { url: imgUrl, path: imagePath });
             if (typeof callbackOriginal === 'function') {
-                await callbackOriginal(prompt, imagePath, generationType, negativePromptPrefix, initiator, prefixedPrompt);
+                await callbackOriginal(prompt, imagePath, generationType, negativePromptPrefix, initiator, prefixedPrompt, format);
             } else {
-                await sendMessage(prompt, imagePath, generationType, negativePromptPrefix, initiator, prefixedPrompt);
+                await sendMessage(prompt, imagePath, generationType, negativePromptPrefix, initiator, prefixedPrompt, format);
+            }
         };
+    }

     const filename = `${characterName}_${humanizedDateTime()}`;
     const base64Image = await saveBase64AsFile(result.data, characterName, filename, result.format);
     callback
-        ? await callback(prompt, base64Image, generationType, additionalNegativePrefix, initiator, prefixedPrompt)
+        ? await callback(prompt, base64Image, generationType, additionalNegativePrefix, initiator, prefixedPrompt, result.format)
-        : await sendMessage(prompt, base64Image, generationType, additionalNegativePrefix, initiator, prefixedPrompt);
+        : await sendMessage(prompt, base64Image, generationType, additionalNegativePrefix, initiator, prefixedPrompt, result.format);
     return base64Image;
+}

@@ -3524,7 +3525,8 @@ async function generateComfyImage(prompt, negativePrompt, signal) {
3524	const text = await promptResult.text();	3525	const text = await promptResult.text();
3525	throw new Error(text);	3526	throw new Error(text);
3526	}	3527	}
3527	return { format: 'png', data: await promptResult.text() };	3528	const { format, data } = await promptResult.json();
		3529	return { format, data };
3528	}	3530	}
3529		3531
3530		3532

  * @param {string} additionalNegativePrefix Additional negative prompt used for the image generation
  * @param {string} initiator The initiator of the image generation
  * @param {string} prefixedPrompt Prompt with an attached specific prefix
+ * @param {string} format Format of the image (e.g., 'png', 'jpg')
  */
-async function sendMessage(prompt, image, generationType, additionalNegativePrefix, initiator, prefixedPrompt) {
+async function sendMessage(prompt, image, generationType, additionalNegativePrefix, initiator, prefixedPrompt, format) {
     const context = getContext();
     const name = context.groupId ? systemUserName : context.name2;
     const template = extension_settings.sd.prompts[generationMode.MESSAGE] || '{{prompt}}';

             image_swipes: [image],
         },
     };
+    if (isVideo(format)) {
+        message.extra.video = image;
+        delete message.extra.image;
+        delete message.extra.image_swipes;
+        delete message.extra.inline_image;
+    }
     context.chat.push(message);
     const messageId = context.chat.length - 1;
     await eventSource.emit(event_types.MESSAGE_RECEIVED, messageId, 'extension');

@@ -4068,7 +4077,7 @@ async function sdMessageButton(e) {
4068	}	4077	}
4069	}	4078	}
4070		4079
4071	function saveGeneratedImage(prompt, image, generationType, negative) {	4080	function saveGeneratedImage(prompt, image, generationType, negative, _initiator, _prefixedPrompt, format) {
4072	// Some message sources may not create the extra object	4081	// Some message sources may not create the extra object
4073	if (typeof message.extra !== 'object' \|\| message.extra === null) {	4082	if (typeof message.extra !== 'object' \|\| message.extra === null) {
4074	message.extra = {};	4083	message.extra = {};

             swipes.push(message.extra.image);
+        }
-        swipes.push(image);
+        const isVideoFormat = isVideo(format);
+        if (isVideoFormat) {
+            message.extra.video = image;
+        } else {
+            swipes.push(image);
+            // If already contains an image and it's not inline - leave it as is
+            message.extra.inline_image = !(message.extra.image && !message.extra.inline_image);
+            message.extra.image = image;
+        }
-        // If already contains an image and it's not inline - leave it as is
-        message.extra.inline_image = !(message.extra.image && !message.extra.inline_image);
-        message.extra.image = image;
         message.extra.title = prompt;
         message.extra.generationType = generationType;
         message.extra.negative = negative;
         appendMediaToMessage(message, $mes);
-        context.saveChat();
+        return context.saveChat();
+    }
+}

src/endpoints/stable-diffusion.js+9 -3 Viewed

         models.forEach(it => it.text = it.text.replace(/\.[^.]*$/, '').replace(/_/g, ' '));
         return response.send(models);
-    } catch (error)     {
+    } catch (error) {
         console.error(error);
         return response.sendStatus(500);
+    }

                 .join('\n') || '';
             throw new Error(`ComfyUI generation did not succeed.\n\n${errorMessages}`.trim());
+        }
-        const imgInfo = Object.keys(item.outputs).map(it => item.outputs[it].images).flat()[0];
+        const outputs = Object.keys(item.outputs).map(it => item.outputs[it]);
+        console.debug('ComfyUI outputs:', outputs);
+        const imgInfo = outputs.map(it => it.images).flat()[0] ?? outputs.map(it => it.gifs).flat()[0];
+        if (!imgInfo) {
+            throw new Error('ComfyUI did not return any recognizable outputs.');
+        }
         const imgUrl = new URL(urlJoin(request.body.url, '/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 format = path.extname(imgInfo.filename).slice(1).toLowerCase() || 'png';
         const imgBuffer = await imgResponse.arrayBuffer();
-        return response.send(Buffer.from(imgBuffer).toString('base64'));
+        return response.send({ format: format, data: Buffer.from(imgBuffer).toString('base64') });
     } catch (error) {
         console.error('ComfyUI error:', error);
         response.status(500).send(error.message);