Skip to content

Fixes *WS node previews #1438

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jtreminio wants to merge 2 commits into
base: master
Choose a base branch
from
+36 −23
Open

Fixes *WS node previews #1438

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
47479fd
Fixes *WS node previews
jtreminio Jun 29, 2026
ce0df75
SwarmSaveAnimatedWebpWS should not save to disk
jtreminio Jun 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 10 additions & 3 deletions src/BuiltinExtensions/ComfyUIBackend/ComfyUIAPIAbstractBackend.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -349,6 +349,7 @@ void yieldProgressUpdate()
bool isExpectingVideo = false;
bool isExpectingText = false;
string currentNode = "";
HashSet<string> 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();
Expand Down Expand Up @@ -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
{
Expand All @@ -540,7 +542,7 @@ async Task doInterruptNow()
JObject historyOut = await SendGet<JObject>($"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)
{
Expand Down Expand Up @@ -639,7 +641,7 @@ public static (MediaType, int, int, int) ComfyRawWebsocketOutputToFormatLabel(by
}
}

private async Task<MediaFile[]> GetAllImagesForHistory(JToken output, T2IParamInput userInput, CancellationToken interrupt)
private async Task<MediaFile[]> GetAllImagesForHistory(JToken output, T2IParamInput userInput, CancellationToken interrupt, HashSet<string> websocketOutputNodes = null)
{
if (Logs.MinimumLevel <= Logs.LogLevel.Verbose)
{
Expand Down Expand Up @@ -675,8 +677,13 @@ private async Task<MediaFile[]> GetAllImagesForHistory(JToken output, T2IParamIn
}
List<MediaFile> outputs = [];
List<string> outputFailures = [];
foreach (JToken outData in output["outputs"].Values())
foreach (JProperty outputProp in output["outputs"].Children<JProperty>())
{
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()}");
Expand Down
18 changes: 10 additions & 8 deletions src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmSaveImageWS.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {}
Expand Down Expand Up @@ -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)
Expand Down
28 changes: 16 additions & 12 deletions src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyExtra/SwarmSaveAnimationWS.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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:
Expand All @@ -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)
Expand All @@ -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()
Expand All @@ -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)

Expand All @@ -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):
Expand Down
Loading