From 47479fd1c81078764c186fbb5809c0d99b6a22f2 Mon Sep 17 00:00:00 2001 From: Juan Treminio Date: Mon, 29 Jun 2026 09:39:38 -0600 Subject: [PATCH 1/2] Fixes *WS node previews --- .../ComfyUIAPIAbstractBackend.cs | 13 ++++++-- .../SwarmComfyCommon/SwarmSaveImageWS.py | 32 ++++++++++++------- .../SwarmComfyExtra/SwarmSaveAnimationWS.py | 28 +++++++++------- 3 files changed, 47 insertions(+), 26 deletions(-) diff --git a/src/BuiltinExtensions/ComfyUIBackend/ComfyUIAPIAbstractBackend.cs b/src/BuiltinExtensions/ComfyUIBackend/ComfyUIAPIAbstractBackend.cs index ff883adb1..b072c89a4 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ComfyUIAPIAbstractBackend.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/ComfyUIAPIAbstractBackend.cs @@ -349,6 +349,7 @@ void yieldProgressUpdate() bool isExpectingVideo = false; bool isExpectingText = false; string currentNode = ""; + HashSet websocketOutputNodes = []; bool isMe = false; // autoCanceller will be cancelled via the using to end the task and not leave it waiting when the method clears using CancellationTokenSource autoCanceller = new(); @@ -517,6 +518,7 @@ async Task doInterruptNow() }; } takeOutput(new T2IEngine.ImageOutput() { File = new Image(output[preBytes..], mediaType), IsReal = isReal, BackendInternalHint = currentNode, GenTimeMS = firstStep == 0 ? -1 : (Environment.TickCount64 - firstStep) }); + websocketOutputNodes.Add(currentNode); } else { @@ -540,7 +542,7 @@ async Task doInterruptNow() JObject historyOut = await SendGet($"history/{promptId}"); if (!historyOut.Properties().IsEmpty()) { - foreach (MediaFile file in await GetAllImagesForHistory(historyOut[promptId], user_input, interrupt)) + foreach (MediaFile file in await GetAllImagesForHistory(historyOut[promptId], user_input, interrupt, websocketOutputNodes)) { if (Program.ServerSettings.AddDebugData) { @@ -639,7 +641,7 @@ public static (MediaType, int, int, int) ComfyRawWebsocketOutputToFormatLabel(by } } - private async Task GetAllImagesForHistory(JToken output, T2IParamInput userInput, CancellationToken interrupt) + private async Task GetAllImagesForHistory(JToken output, T2IParamInput userInput, CancellationToken interrupt, HashSet websocketOutputNodes = null) { if (Logs.MinimumLevel <= Logs.LogLevel.Verbose) { @@ -675,8 +677,13 @@ private async Task GetAllImagesForHistory(JToken output, T2IParamIn } List outputs = []; List outputFailures = []; - foreach (JToken outData in output["outputs"].Values()) + foreach (JProperty outputProp in output["outputs"].Children()) { + if (websocketOutputNodes is not null && websocketOutputNodes.Contains(outputProp.Name)) + { + continue; + } + JToken outData = outputProp.Value; if (outData is null) { Logs.Debug($"null output data from ComfyUI server: {output.ToDenseDebugString()}"); diff --git a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py index 073f32512..9d2d15425 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py +++ b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py @@ -2,7 +2,8 @@ import numpy as np import comfy.utils from server import PromptServer, BinaryEventTypes -import time, io, struct +import time, io, struct, os +import folder_paths SPECIAL_ID = 12345 # Tells swarm that the node is going to output final images VIDEO_ID = 12346 @@ -38,26 +39,24 @@ 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) + pbar = comfy.utils.ProgressBar(images.shape[0]) step = 0 for image in images: + i = 255.0 * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + pbar.update_absolute(step, images.shape[0], ("PNG", img, None)) if bit_depth == "raw": - i = 255.0 * image.cpu().numpy() - img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) def do_save(out): img.save(out, format='BMP') send_image_to_server_raw(1, do_save, SPECIAL_ID, event_type=10) elif bit_depth == "16bit": i = 65535.0 * image.cpu().numpy() - img = self.convert_img_16bit(np.clip(i, 0, 65535).astype(np.uint16)) - send_image_to_server_raw(2, lambda out: out.write(img), SPECIAL_ID) + img16 = self.convert_img_16bit(np.clip(i, 0, 65535).astype(np.uint16)) + send_image_to_server_raw(2, lambda out: out.write(img16), SPECIAL_ID) else: - i = 255.0 * image.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(step, SPECIAL_ID, ("PNG", img, None)) step += 1 return {} @@ -104,17 +103,28 @@ def INPUT_TYPES(s): def save_images(self, images, fps, lossless, quality, method): method = self.methods.get(method) + if images.shape[0] == 0: + return { } pil_images = [] for image in images: i = 255. * image.cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) pil_images.append(img) + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path("swarm_preview_", folder_paths.get_temp_directory(), images[0].shape[1], images[0].shape[0]) + file = f"{filename}_{counter:05}_.webp" + file_path = os.path.join(full_output_folder, file) + pil_images[0].save(file_path, save_all=True, duration=int(1000.0/fps), append_images=pil_images[1 : len(pil_images)], lossless=lossless, quality=quality, method=method, format='WEBP') + def do_save(out): - pil_images[0].save(out, save_all=True, duration=int(1000.0/fps), append_images=pil_images[1 : len(pil_images)], lossless=lossless, quality=quality, method=method, format='WEBP') + with open(file_path, "rb") as f: + out.write(f.read()) send_image_to_server_raw(3, do_save, VIDEO_ID) - return { } + ui = { "images": [{ "filename": file, "subfolder": subfolder, "type": "temp" }] } + if len(pil_images) > 1: + ui["animated"] = (True,) + return { "ui": ui } @classmethod def IS_CHANGED(s, images, fps, lossless, quality, method): diff --git a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyExtra/SwarmSaveAnimationWS.py b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyExtra/SwarmSaveAnimationWS.py index bc9ea0714..4e9259a8c 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 comfy, folder_paths, io, struct, subprocess, os, sys, time, wave from PIL import Image import numpy as np from server import PromptServer, BinaryEventTypes @@ -49,16 +49,18 @@ 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) + pbar = comfy.utils.ProgressBar(images.shape[0]) i = 255.0 * images[0].cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + pbar.update_absolute(0, images.shape[0], ("PNG", img, None)) 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 { } + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path("swarm_preview_", folder_paths.get_temp_directory(), images[0].shape[1], images[0].shape[0]) out_img = io.BytesIO() + file = None if format in ["webp", "gif"]: if format == "webp": type_num = 3 @@ -69,7 +71,11 @@ def do_save(out): i = 255. * image.cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) pil_images.append(img) - pil_images[0].save(out_img, save_all=True, duration=int(1000.0 / fps), append_images=pil_images[1 : len(pil_images)], lossless=lossless, quality=quality, method=method, format=format.upper(), loop=0) + file = f"{filename}_{counter:05}_.{format}" + file_path = os.path.join(full_output_folder, file) + pil_images[0].save(file_path, save_all=True, duration=int(1000.0 / fps), append_images=pil_images[1 : len(pil_images)], lossless=lossless, quality=quality, method=method, format=format.upper(), loop=0) + with open(file_path, "rb") as f: + out_img.write(f.read()) else: i = 255. * images.cpu().numpy() raw_images = np.clip(i, 0, 255).astype(np.uint8) @@ -101,9 +107,8 @@ def do_save(out): video_args = ["-filter_complex", "split=2 [a][b]; [a] palettegen [pal]; [b] [pal] paletteuse"] ext = "gif" type_num = 4 - path = folder_paths.get_save_image_path("swarm_tmp_", folder_paths.get_temp_directory())[0] - rand = '%016x' % random.getrandbits(64) - file = os.path.join(path, f"swarm_tmp_{rand}.{ext}") + file = f"{filename}_{counter:05}_.{ext}" + file_path = os.path.join(full_output_folder, file) file_2 = None audio_input = [] if audio is not None and audio_args is not None: @@ -123,7 +128,7 @@ def do_save(out): audio_np = np.concatenate([audio_np, padding], axis=1) audio_np = audio_np.T audio_int16 = (np.clip(audio_np, -1.0, 1.0) * 32767).astype(np.int16) - file_2 = os.path.join(path, f"swarm_tmp_{rand}_audio.wav") + file_2 = os.path.join(full_output_folder, f"{filename}_{counter:05}_audio.wav") with wave.open(file_2, 'wb') as wav_file: wav_file.setnchannels(channels) wav_file.setsampwidth(2) @@ -132,7 +137,7 @@ def do_save(out): audio_input = ["-i", file_2] else: audio_args = [] - result = subprocess.run(args + audio_input + video_args + audio_args + [file], input=raw_images.tobytes(), capture_output=True) + result = subprocess.run(args + audio_input + video_args + audio_args + [file_path], input=raw_images.tobytes(), capture_output=True) if result.returncode != 0: print(f"ffmpeg failed with return code {result.returncode}", file=sys.stderr) f_out = result.stdout.decode("utf-8").strip() @@ -143,9 +148,8 @@ def do_save(out): print("ffmpeg error: " + f_err, file=sys.stderr) raise Exception(f"ffmpeg failed: {f_err}") # TODO: Is there a way to get ffmpeg to operate entirely in memory? - with open(file, "rb") as f: + with open(file_path, "rb") as f: out_img.write(f.read()) - os.remove(file) if file_2 is not None: os.remove(file_2) @@ -159,7 +163,7 @@ def do_save(out): server.send_sync("progress", {"value": 12346, "max": 12346}, sid=server.client_id) server.send_sync(BinaryEventTypes.PREVIEW_IMAGE, preview_bytes, sid=server.client_id) - return { } + return { "ui": { "images": [{ "filename": file, "subfolder": subfolder, "type": "temp" }], "animated": (True,) } } @classmethod def IS_CHANGED(s, images, fps, lossless, quality, method, format, audio=None): From ce0df75202dba4c7dd19bdc45c615b6f8a7a1ccf Mon Sep 17 00:00:00 2001 From: Juan Treminio Date: Mon, 29 Jun 2026 18:47:57 -0600 Subject: [PATCH 2/2] SwarmSaveAnimatedWebpWS should not save to disk --- .../SwarmComfyCommon/SwarmSaveImageWS.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py index 9d2d15425..f93acbc9c 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py +++ b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py @@ -2,8 +2,7 @@ import numpy as np import comfy.utils from server import PromptServer, BinaryEventTypes -import time, io, struct, os -import folder_paths +import time, io, struct SPECIAL_ID = 12345 # Tells swarm that the node is going to output final images VIDEO_ID = 12346 @@ -111,20 +110,13 @@ def save_images(self, images, fps, lossless, quality, method): img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) pil_images.append(img) - full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path("swarm_preview_", folder_paths.get_temp_directory(), images[0].shape[1], images[0].shape[0]) - file = f"{filename}_{counter:05}_.webp" - file_path = os.path.join(full_output_folder, file) - pil_images[0].save(file_path, save_all=True, duration=int(1000.0/fps), append_images=pil_images[1 : len(pil_images)], lossless=lossless, quality=quality, method=method, format='WEBP') + comfy.utils.ProgressBar(1).update_absolute(0, 1, ("PNG", pil_images[0], None)) def do_save(out): - with open(file_path, "rb") as f: - out.write(f.read()) + pil_images[0].save(out, save_all=True, duration=int(1000.0/fps), append_images=pil_images[1 : len(pil_images)], lossless=lossless, quality=quality, method=method, format='WEBP') send_image_to_server_raw(3, do_save, VIDEO_ID) - ui = { "images": [{ "filename": file, "subfolder": subfolder, "type": "temp" }] } - if len(pil_images) > 1: - ui["animated"] = (True,) - return { "ui": ui } + return { } @classmethod def IS_CHANGED(s, images, fps, lossless, quality, method):