46 lines
1.2 KiB
Python
46 lines
1.2 KiB
Python
from PIL import Image, ImageOps
|
|
from io import BytesIO
|
|
import numpy as np
|
|
import struct
|
|
import comfy.utils
|
|
import time
|
|
|
|
#You can use this node to save full size images through the websocket, the
|
|
#images will be sent in exactly the same format as the image previews: as
|
|
#binary images on the websocket with a 8 byte header indicating the type
|
|
#of binary message (first 4 bytes) and the image format (next 4 bytes).
|
|
|
|
#Note that no metadata will be put in the images saved with this node.
|
|
|
|
class SaveImageWebsocket:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":
|
|
{"images": ("IMAGE", ),}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "save_images"
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
CATEGORY = "api/image"
|
|
|
|
def save_images(self, images):
|
|
pbar = comfy.utils.ProgressBar(images.shape[0])
|
|
step = 0
|
|
for image in images:
|
|
i = 255. * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
pbar.update_absolute(step, images.shape[0], ("PNG", img, None))
|
|
step += 1
|
|
|
|
return {}
|
|
|
|
def IS_CHANGED(s, images):
|
|
return time.time()
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SaveImageWebsocket": SaveImageWebsocket,
|
|
}
|