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..f93acbc9c 100644 --- a/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py +++ b/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py @@ -38,26 +38,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,12 +102,16 @@ 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) + 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) 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):