File size: 5,041 Bytes
efc1c32 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
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}")
# TODO: Is there a way to get ffmpeg to operate entirely in memory?
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,
}
|