Fixes *WS node previews

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 {}
 

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/'));
+});

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