Skip to content

Commit 7f89cb4

Browse files
Add a disabled SaveImageWebsocket custom node.
This node can be used to efficiently get images without saving them to disk when using ComfyUI as a backend.
1 parent 38b7ac6 commit 7f89cb4

File tree

1 file changed

+49
-0
lines changed

1 file changed

+49
-0
lines changed
Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
from PIL import Image, ImageOps
2+
from io import BytesIO
3+
import numpy as np
4+
import struct
5+
import comfy.utils
6+
import time
7+
8+
#You can use this node to save full size images through the websocket, the
9+
#images will be sent in exactly the same format as the image previews: as
10+
#binary images on the websocket with a 8 byte header indicating the type
11+
#of binary message (first 4 bytes) and the image format (next 4 bytes).
12+
13+
#The reason this node is disabled by default is because there is a small
14+
#issue when using it with the default ComfyUI web interface: When generating
15+
#batches only the last image will be shown in the UI.
16+
17+
#Note that no metadata will be put in the images saved with this node.
18+
19+
class SaveImageWebsocket:
20+
@classmethod
21+
def INPUT_TYPES(s):
22+
return {"required":
23+
{"images": ("IMAGE", ),}
24+
}
25+
26+
RETURN_TYPES = ()
27+
FUNCTION = "save_images"
28+
29+
OUTPUT_NODE = True
30+
31+
CATEGORY = "image"
32+
33+
def save_images(self, images):
34+
pbar = comfy.utils.ProgressBar(images.shape[0])
35+
step = 0
36+
for image in images:
37+
i = 255. * image.cpu().numpy()
38+
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
39+
pbar.update_absolute(step, images.shape[0], ("PNG", img, None))
40+
step += 1
41+
42+
return {}
43+
44+
def IS_CHANGED(s, images):
45+
return time.time()
46+
47+
NODE_CLASS_MAPPINGS = {
48+
"SaveImageWebsocket": SaveImageWebsocket,
49+
}

0 commit comments

Comments
 (0)