diff --git a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py index 073f32512..d47e95d1f 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py +++ b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py @@ -1,6 +1,5 @@ from PIL import Image import numpy as np -import comfy.utils from server import PromptServer, BinaryEventTypes import time, io, struct @@ -38,8 +37,6 @@ def INPUT_TYPES(s): DESCRIPTION = "Acts like a special version of 'SaveImage' that doesn't actual save to disk, instead it sends directly over websocket. This is intended so that SwarmUI can save the image itself rather than having Comfy's Core save it." def save_images(self, images, bit_depth = "8bit"): - pbar = comfy.utils.ProgressBar(SPECIAL_ID) - step = 0 for image in images: if bit_depth == "raw": i = 255.0 * image.cpu().numpy() @@ -57,8 +54,6 @@ def do_save(out): def do_save(out): img.save(out, format='PNG') send_image_to_server_raw(2, do_save, SPECIAL_ID) - #pbar.update_absolute(step, SPECIAL_ID, ("PNG", img, None)) - step += 1 return {} diff --git a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/web/swarmhelper.js b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/web/swarmhelper.js index 94a00e848..af9f3f599 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/web/swarmhelper.js +++ b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/web/swarmhelper.js @@ -1,3 +1,72 @@ import { api } from '../../scripts/api.js'; +import { app } from '../../scripts/app.js'; window.swarmApiDirect = api; + +let swarmSaveNodes = ['SwarmSaveImageWS', 'SwarmSaveAnimatedWebpWS', 'SwarmSaveAnimationWS']; +let swarmExecutingNode = null; + +function swarmSniffMime(bytes) { + let ascii = (start, len) => String.fromCharCode(...bytes.slice(start, start + len)); + if (ascii(0, 4) == 'RIFF' && ascii(8, 4) == 'WEBP') { + return 'image/webp'; + } + if (ascii(0, 3) == 'GIF') { + return 'image/gif'; + } + if (ascii(4, 4) == 'ftyp') { + return 'video/mp4'; + } + if (bytes[0] == 0x1A && bytes[1] == 0x45 && bytes[2] == 0xDF && bytes[3] == 0xA3) { + return 'video/webm'; + } + return null; +} + +function swarmShowPreview(node, blob, isVideo) { + let widget = node.widgets?.find(w => w.name == 'swarm_ws_preview'); + if (!widget) { + let container = document.createElement('div'); + container.classList.add('comfy-img-preview'); + widget = node.addDOMWidget('swarm_ws_preview', 'swarm_ws_preview', container, { canvasOnly: true, hideOnZoom: false }); + widget.serialize = false; + widget.computeLayoutSize = () => ({ minWidth: 0, minHeight: widget.swarmMinHeight || 256 }); + } + let element = document.createElement(isVideo ? 'video' : 'img'); + if (isVideo) { + element.muted = true; + element.autoplay = true; + element.loop = true; + element.playsInline = true; + } + element.style.width = '100%'; + element.style.height = '100%'; + element.style.objectFit = 'contain'; + element.onload = element.onloadedmetadata = () => { + let width = element.videoWidth || element.naturalWidth; + let height = element.videoHeight || element.naturalHeight; + widget.swarmMinHeight = height * Math.min(1, (node.size[0] || 256) / width); + node.graph?.setDirtyCanvas(true); + }; + let old = widget.element.firstChild; + if (old) { + URL.revokeObjectURL(old.src); + } + element.src = URL.createObjectURL(blob); + widget.element.replaceChildren(element); + delete app.nodePreviewImages?.[node.id]; +} + +api.addEventListener('executing', ({ detail }) => { + swarmExecutingNode = detail?.display_node ?? detail?.node ?? detail; +}); + +api.addEventListener('b_preview', async ({ detail }) => { + let node = app.graph.getNodeById(swarmExecutingNode); + if (!node || !swarmSaveNodes.includes(node.comfyClass)) { + return; + } + let head = new Uint8Array(await detail.slice(0, 12).arrayBuffer()); + let mime = swarmSniffMime(head) ?? detail.type; + swarmShowPreview(node, detail.slice(0, detail.size, mime), mime.startsWith('video/')); +}); diff --git a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyExtra/SwarmSaveAnimationWS.py b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyExtra/SwarmSaveAnimationWS.py index bc9ea0714..f76f2c3e1 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyExtra/SwarmSaveAnimationWS.py +++ b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyExtra/SwarmSaveAnimationWS.py @@ -1,4 +1,4 @@ -import comfy, folder_paths, io, struct, subprocess, os, random, sys, time, wave +import folder_paths, io, struct, subprocess, os, random, sys, time, wave from PIL import Image import numpy as np from server import PromptServer, BinaryEventTypes @@ -49,13 +49,11 @@ def save_images(self, images, fps, lossless, quality, method, format, audio=None if images.shape[0] == 0: return { } if images.shape[0] == 1: - pbar = comfy.utils.ProgressBar(SPECIAL_ID) i = 255.0 * images[0].cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) def do_save(out): img.save(out, format='PNG') send_image_to_server_raw(2, do_save, SPECIAL_ID) - #pbar.update_absolute(0, SPECIAL_ID, ("PNG", img, None)) return { } out_img = io.BytesIO()