| import comfy, folder_paths, io, struct, subprocess, os, random, sys, time
|
| from PIL import Image
|
| import numpy as np
|
| from server import PromptServer, BinaryEventTypes
|
| from imageio_ffmpeg import get_ffmpeg_exe
|
|
|
| SPECIAL_ID = 12345
|
| VIDEO_ID = 12346
|
| FFMPEG_PATH = get_ffmpeg_exe()
|
|
|
|
|
| class SwarmSaveAnimationWS:
|
| methods = {"default": 4, "fastest": 0, "slowest": 6}
|
|
|
| @classmethod
|
| def INPUT_TYPES(s):
|
| return {
|
| "required": {
|
| "images": ("IMAGE", ),
|
| "fps": ("FLOAT", {"default": 6.0, "min": 0.01, "max": 1000.0, "step": 0.01}),
|
| "lossless": ("BOOLEAN", {"default": True}),
|
| "quality": ("INT", {"default": 80, "min": 0, "max": 100}),
|
| "method": (list(s.methods.keys()),),
|
| "format": (["webp", "gif", "gif-hd", "h264-mp4", "h265-mp4", "webm", "prores"],),
|
| },
|
| }
|
|
|
| CATEGORY = "SwarmUI/video"
|
| RETURN_TYPES = ()
|
| FUNCTION = "save_images"
|
| OUTPUT_NODE = True
|
|
|
| def save_images(self, images, fps, lossless, quality, method, format):
|
| method = self.methods.get(method)
|
| if images.shape[0] == 0:
|
| return { }
|
| if images.shape[0] == 1:
|
| pbar = comfy.utils.ProgressBar(SPECIAL_ID)
|
| i = 255.0 * images[0].cpu().numpy()
|
| img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
| pbar.update_absolute(0, SPECIAL_ID, ("PNG", img, None))
|
| return { }
|
|
|
| out_img = io.BytesIO()
|
| if format in ["webp", "gif"]:
|
| if format == "webp":
|
| type_num = 3
|
| else:
|
| type_num = 4
|
| 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)
|
| 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)
|
| else:
|
| i = 255. * images.cpu().numpy()
|
| raw_images = np.clip(i, 0, 255).astype(np.uint8)
|
| args = [FFMPEG_PATH, "-v", "error", "-f", "rawvideo", "-pix_fmt", "rgb24",
|
| "-s", f"{len(raw_images[0][0])}x{len(raw_images[0])}", "-r", str(fps), "-i", "-", "-n" ]
|
| if format == "h264-mp4":
|
| args += ["-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "19"]
|
| ext = "mp4"
|
| type_num = 5
|
| elif format == "h265-mp4":
|
| args += ["-c:v", "libx265", "-pix_fmt", "yuv420p"]
|
| ext = "mp4"
|
| type_num = 5
|
| elif format == "webm":
|
| args += ["-pix_fmt", "yuv420p", "-crf", "23"]
|
| ext = "webm"
|
| type_num = 6
|
| elif format == "prores":
|
| args += ["-c:v", "prores_ks", "-profile:v", "3", "-pix_fmt", "yuv422p10le"]
|
| ext = "mov"
|
| type_num = 7
|
| elif format == "gif-hd":
|
| 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}")
|
| result = subprocess.run(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()
|
| f_err = result.stderr.decode("utf-8").strip()
|
| if f_out:
|
| print("ffmpeg out: " + f_out, file=sys.stderr)
|
| if f_err:
|
| print("ffmpeg error: " + f_err, file=sys.stderr)
|
| raise Exception(f"ffmpeg failed: {f_err}")
|
|
|
| with open(file, "rb") as f:
|
| out_img.write(f.read())
|
| os.remove(file)
|
|
|
| out = io.BytesIO()
|
| header = struct.pack(">I", type_num)
|
| out.write(header)
|
| out.write(out_img.getvalue())
|
| out.seek(0)
|
| preview_bytes = out.getvalue()
|
| server = PromptServer.instance
|
| 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 { }
|
|
|
| @classmethod
|
| def IS_CHANGED(s, images, fps, lossless, quality, method, format):
|
| return time.time()
|
|
|
|
|
| NODE_CLASS_MAPPINGS = {
|
| "SwarmSaveAnimationWS": SwarmSaveAnimationWS,
|
| }
|
|
|
