From 47479fd1c81078764c186fbb5809c0d99b6a22f2 Mon Sep 17 00:00:00 2001 From: Juan Treminio Date: Mon, 29 Jun 2026 09:39:38 -0600 Subject: [PATCH 1/4] 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/4] 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): From 9437fbc22498e33c1117ee08b3b34b508caf55cc Mon Sep 17 00:00:00 2001 From: Juan Treminio Date: Wed, 15 Jul 2026 13:28:31 -0600 Subject: [PATCH 3/4] tidy up WS preview handling --- .../ComfyUIAPIAbstractBackend.cs | 17 +++--- .../SwarmComfyCommon/SwarmSaveImageWS.py | 52 +++++++++++-------- .../SwarmComfyExtra/SwarmSaveAnimationWS.py | 26 +++++----- 3 files changed, 51 insertions(+), 44 deletions(-) diff --git a/src/BuiltinExtensions/ComfyUIBackend/ComfyUIAPIAbstractBackend.cs b/src/BuiltinExtensions/ComfyUIBackend/ComfyUIAPIAbstractBackend.cs index b072c89a4..1e6937590 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ComfyUIAPIAbstractBackend.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/ComfyUIAPIAbstractBackend.cs @@ -349,7 +349,6 @@ 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(); @@ -518,7 +517,6 @@ 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 { @@ -542,7 +540,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, websocketOutputNodes)) + foreach (MediaFile file in await GetAllImagesForHistory(historyOut[promptId], user_input, interrupt)) { if (Program.ServerSettings.AddDebugData) { @@ -641,7 +639,7 @@ public static (MediaType, int, int, int) ComfyRawWebsocketOutputToFormatLabel(by } } - private async Task GetAllImagesForHistory(JToken output, T2IParamInput userInput, CancellationToken interrupt, HashSet websocketOutputNodes = null) + private async Task GetAllImagesForHistory(JToken output, T2IParamInput userInput, CancellationToken interrupt) { if (Logs.MinimumLevel <= Logs.LogLevel.Verbose) { @@ -677,13 +675,8 @@ private async Task GetAllImagesForHistory(JToken output, T2IParamIn } List outputs = []; List outputFailures = []; - foreach (JProperty outputProp in output["outputs"].Children()) + foreach (JToken outData in output["outputs"].Values()) { - 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()}"); @@ -697,6 +690,10 @@ async Task LoadOutput(JObject outImage, MediaMetaType metaType) if ($"{outImage["type"]}" == "temp") { imType = "temp"; + if (fname.StartsWith("swarm_preview_")) + { + return; + } } string url = $"filename={HttpUtility.UrlEncode(fname)}&type={imType}"; if (outImage.TryGetValue("subfolder", out JToken subFolder) && !string.IsNullOrWhiteSpace($"{subFolder}")) diff --git a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py index f93acbc9c..8a09b8e0e 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py +++ b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py @@ -1,12 +1,13 @@ from PIL import Image import numpy as np -import comfy.utils from server import PromptServer, BinaryEventTypes -import time, io, struct +import time, io, struct, os, random +import folder_paths SPECIAL_ID = 12345 # Tells swarm that the node is going to output final images VIDEO_ID = 12346 TEXT_ID = 12347 +PREVIEW_PREFIX = "swarm_preview_" + '%08x' % random.getrandbits(32) def send_image_to_server_raw(type_num: int, save_me: callable, id: int, event_type: int = BinaryEventTypes.PREVIEW_IMAGE): out = io.BytesIO() @@ -38,27 +39,31 @@ 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(images.shape[0]) - step = 0 + if images.shape[0] == 0: + return { } + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(PREVIEW_PREFIX, folder_paths.get_temp_directory(), images[0].shape[1], images[0].shape[0]) + results = [] 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)) + i = image.cpu().numpy() + img = Image.fromarray(np.clip(255.0 * i, 0, 255).astype(np.uint8)) + png = io.BytesIO() + img.save(png, format='PNG') + file = f"{filename}_{counter:05}_.png" + with open(os.path.join(full_output_folder, file), "wb") as f: + f.write(png.getvalue()) + results.append({ "filename": file, "subfolder": subfolder, "type": "temp" }) + counter += 1 if bit_depth == "raw": 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() - img16 = self.convert_img_16bit(np.clip(i, 0, 65535).astype(np.uint16)) + img16 = self.convert_img_16bit(np.clip(65535.0 * i, 0, 65535).astype(np.uint16)) send_image_to_server_raw(2, lambda out: out.write(img16), SPECIAL_ID) else: - def do_save(out): - img.save(out, format='PNG') - send_image_to_server_raw(2, do_save, SPECIAL_ID) - step += 1 + send_image_to_server_raw(2, lambda out: out.write(png.getvalue()), SPECIAL_ID) - return {} + return { "ui": { "images": results } } def convert_img_16bit(self, img_np): try: @@ -110,13 +115,18 @@ def save_images(self, images, fps, lossless, quality, method): img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) pil_images.append(img) - comfy.utils.ProgressBar(1).update_absolute(0, 1, ("PNG", pil_images[0], None)) - - 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') - send_image_to_server_raw(3, do_save, VIDEO_ID) - - return { } + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(PREVIEW_PREFIX, folder_paths.get_temp_directory(), images[0].shape[1], images[0].shape[0]) + file = f"{filename}_{counter:05}_.webp" + webp = io.BytesIO() + pil_images[0].save(webp, 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(os.path.join(full_output_folder, file), "wb") as f: + f.write(webp.getvalue()) + send_image_to_server_raw(3, lambda out: out.write(webp.getvalue()), VIDEO_ID) + + 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 4e9259a8c..2c7d0290b 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, 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 @@ -6,6 +6,7 @@ SPECIAL_ID = 12345 VIDEO_ID = 12346 +PREVIEW_PREFIX = "swarm_preview_" + '%08x' % random.getrandbits(32) FFMPEG_PATH = get_ffmpeg_exe() def send_image_to_server_raw(type_num: int, save_me: callable, id: int, event_type: int = BinaryEventTypes.PREVIEW_IMAGE): @@ -48,19 +49,19 @@ def save_images(self, images, fps, lossless, quality, method, format, audio=None method = self.methods.get(method) if images.shape[0] == 0: return { } + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(PREVIEW_PREFIX, folder_paths.get_temp_directory(), images[0].shape[1], images[0].shape[0]) if images.shape[0] == 1: - 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) - return { } + file = f"{filename}_{counter:05}_.png" + png = io.BytesIO() + img.save(png, format='PNG') + with open(os.path.join(full_output_folder, file), "wb") as f: + f.write(png.getvalue()) + send_image_to_server_raw(2, lambda out: out.write(png.getvalue()), SPECIAL_ID) + return { "ui": { "images": [{ "filename": file, "subfolder": subfolder, "type": "temp" }] } } - 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 @@ -72,10 +73,9 @@ def do_save(out): img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) pil_images.append(img) 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()) + 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) + with open(os.path.join(full_output_folder, file), "wb") as f: + f.write(out_img.getvalue()) else: i = 255. * images.cpu().numpy() raw_images = np.clip(i, 0, 255).astype(np.uint8) From 5b72f05d450469a95fffba6598ba21280f70355f Mon Sep 17 00:00:00 2001 From: Juan Treminio Date: Wed, 15 Jul 2026 22:40:01 -0500 Subject: [PATCH 4/4] fix comfy tab previews for WS save nodes via frontend js --- .../ComfyUIAPIAbstractBackend.cs | 4 -- .../SwarmComfyCommon/SwarmSaveImageWS.py | 51 +++++--------- .../SwarmComfyCommon/web/swarmhelper.js | 69 +++++++++++++++++++ .../SwarmComfyExtra/SwarmSaveAnimationWS.py | 30 ++++---- 4 files changed, 98 insertions(+), 56 deletions(-) diff --git a/src/BuiltinExtensions/ComfyUIBackend/ComfyUIAPIAbstractBackend.cs b/src/BuiltinExtensions/ComfyUIBackend/ComfyUIAPIAbstractBackend.cs index 1e6937590..ff883adb1 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ComfyUIAPIAbstractBackend.cs +++ b/src/BuiltinExtensions/ComfyUIBackend/ComfyUIAPIAbstractBackend.cs @@ -690,10 +690,6 @@ async Task LoadOutput(JObject outImage, MediaMetaType metaType) if ($"{outImage["type"]}" == "temp") { imType = "temp"; - if (fname.StartsWith("swarm_preview_")) - { - return; - } } string url = $"filename={HttpUtility.UrlEncode(fname)}&type={imType}"; if (outImage.TryGetValue("subfolder", out JToken subFolder) && !string.IsNullOrWhiteSpace($"{subFolder}")) diff --git a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py index 8a09b8e0e..d47e95d1f 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py +++ b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py @@ -1,13 +1,11 @@ from PIL import Image import numpy as np from server import PromptServer, BinaryEventTypes -import time, io, struct, os, random -import folder_paths +import time, io, struct SPECIAL_ID = 12345 # Tells swarm that the node is going to output final images VIDEO_ID = 12346 TEXT_ID = 12347 -PREVIEW_PREFIX = "swarm_preview_" + '%08x' % random.getrandbits(32) def send_image_to_server_raw(type_num: int, save_me: callable, id: int, event_type: int = BinaryEventTypes.PREVIEW_IMAGE): out = io.BytesIO() @@ -39,31 +37,25 @@ 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"): - if images.shape[0] == 0: - return { } - full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(PREVIEW_PREFIX, folder_paths.get_temp_directory(), images[0].shape[1], images[0].shape[0]) - results = [] for image in images: - i = image.cpu().numpy() - img = Image.fromarray(np.clip(255.0 * i, 0, 255).astype(np.uint8)) - png = io.BytesIO() - img.save(png, format='PNG') - file = f"{filename}_{counter:05}_.png" - with open(os.path.join(full_output_folder, file), "wb") as f: - f.write(png.getvalue()) - results.append({ "filename": file, "subfolder": subfolder, "type": "temp" }) - counter += 1 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": - img16 = self.convert_img_16bit(np.clip(65535.0 * i, 0, 65535).astype(np.uint16)) - send_image_to_server_raw(2, lambda out: out.write(img16), SPECIAL_ID) + 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) else: - send_image_to_server_raw(2, lambda out: out.write(png.getvalue()), SPECIAL_ID) + 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) - return { "ui": { "images": results } } + return {} def convert_img_16bit(self, img_np): try: @@ -107,26 +99,17 @@ 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(PREVIEW_PREFIX, folder_paths.get_temp_directory(), images[0].shape[1], images[0].shape[0]) - file = f"{filename}_{counter:05}_.webp" - webp = io.BytesIO() - pil_images[0].save(webp, 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(os.path.join(full_output_folder, file), "wb") as f: - f.write(webp.getvalue()) - send_image_to_server_raw(3, lambda out: out.write(webp.getvalue()), VIDEO_ID) - - ui = { "images": [{ "filename": file, "subfolder": subfolder, "type": "temp" }] } - if len(pil_images) > 1: - ui["animated"] = (True,) - return { "ui": ui } + 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') + send_image_to_server_raw(3, do_save, VIDEO_ID) + + return { } @classmethod def IS_CHANGED(s, images, fps, lossless, quality, method): 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 2c7d0290b..f76f2c3e1 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyExtra/SwarmSaveAnimationWS.py +++ b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyExtra/SwarmSaveAnimationWS.py @@ -6,7 +6,6 @@ SPECIAL_ID = 12345 VIDEO_ID = 12346 -PREVIEW_PREFIX = "swarm_preview_" + '%08x' % random.getrandbits(32) FFMPEG_PATH = get_ffmpeg_exe() def send_image_to_server_raw(type_num: int, save_me: callable, id: int, event_type: int = BinaryEventTypes.PREVIEW_IMAGE): @@ -49,17 +48,13 @@ def save_images(self, images, fps, lossless, quality, method, format, audio=None method = self.methods.get(method) if images.shape[0] == 0: return { } - full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(PREVIEW_PREFIX, folder_paths.get_temp_directory(), images[0].shape[1], images[0].shape[0]) if images.shape[0] == 1: i = 255.0 * images[0].cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) - file = f"{filename}_{counter:05}_.png" - png = io.BytesIO() - img.save(png, format='PNG') - with open(os.path.join(full_output_folder, file), "wb") as f: - f.write(png.getvalue()) - send_image_to_server_raw(2, lambda out: out.write(png.getvalue()), SPECIAL_ID) - return { "ui": { "images": [{ "filename": file, "subfolder": subfolder, "type": "temp" }] } } + def do_save(out): + img.save(out, format='PNG') + send_image_to_server_raw(2, do_save, SPECIAL_ID) + return { } out_img = io.BytesIO() if format in ["webp", "gif"]: @@ -72,10 +67,7 @@ def save_images(self, images, fps, lossless, quality, method, format, audio=None i = 255. * image.cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) pil_images.append(img) - file = f"{filename}_{counter:05}_.{format}" 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) - with open(os.path.join(full_output_folder, file), "wb") as f: - f.write(out_img.getvalue()) else: i = 255. * images.cpu().numpy() raw_images = np.clip(i, 0, 255).astype(np.uint8) @@ -107,8 +99,9 @@ def save_images(self, images, fps, lossless, quality, method, format, audio=None video_args = ["-filter_complex", "split=2 [a][b]; [a] palettegen [pal]; [b] [pal] paletteuse"] ext = "gif" type_num = 4 - file = f"{filename}_{counter:05}_.{ext}" - file_path = os.path.join(full_output_folder, file) + 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_2 = None audio_input = [] if audio is not None and audio_args is not None: @@ -128,7 +121,7 @@ def save_images(self, images, fps, lossless, quality, method, format, audio=None 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(full_output_folder, f"{filename}_{counter:05}_audio.wav") + file_2 = os.path.join(path, f"swarm_tmp_{rand}_audio.wav") with wave.open(file_2, 'wb') as wav_file: wav_file.setnchannels(channels) wav_file.setsampwidth(2) @@ -137,7 +130,7 @@ def save_images(self, images, fps, lossless, quality, method, format, audio=None audio_input = ["-i", file_2] else: audio_args = [] - result = subprocess.run(args + audio_input + video_args + audio_args + [file_path], input=raw_images.tobytes(), capture_output=True) + result = subprocess.run(args + audio_input + video_args + audio_args + [file], 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() @@ -148,8 +141,9 @@ def save_images(self, images, fps, lossless, quality, method, format, audio=None 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_path, "rb") as f: + with open(file, "rb") as f: out_img.write(f.read()) + os.remove(file) if file_2 is not None: os.remove(file_2) @@ -163,7 +157,7 @@ def save_images(self, images, fps, lossless, quality, method, format, audio=None 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 { "ui": { "images": [{ "filename": file, "subfolder": subfolder, "type": "temp" }], "animated": (True,) } } + return { } @classmethod def IS_CHANGED(s, images, fps, lossless, quality, method, format, audio=None):