diff --git a/comfy_extras/nodes_video.py b/comfy_extras/nodes_video.py index 5c096c232..d4d0dadff 100644 --- a/comfy_extras/nodes_video.py +++ b/comfy_extras/nodes_video.py @@ -39,29 +39,31 @@ class SaveWEBM(io.ComfyNode): file = f"{filename}_{counter:05}_.webm" container = av.open(os.path.join(full_output_folder, file), mode="w") - if cls.hidden.prompt is not None: - container.metadata["prompt"] = json.dumps(cls.hidden.prompt) + try: + if cls.hidden.prompt is not None: + container.metadata["prompt"] = json.dumps(cls.hidden.prompt) - if cls.hidden.extra_pnginfo is not None: - for x in cls.hidden.extra_pnginfo: - container.metadata[x] = json.dumps(cls.hidden.extra_pnginfo[x]) + if cls.hidden.extra_pnginfo is not None: + for x in cls.hidden.extra_pnginfo: + container.metadata[x] = json.dumps(cls.hidden.extra_pnginfo[x]) - codec_map = {"vp9": "libvpx-vp9", "av1": "libsvtav1"} - stream = container.add_stream(codec_map[codec], rate=Fraction(round(fps * 1000), 1000)) - stream.width = images.shape[-2] - stream.height = images.shape[-3] - stream.pix_fmt = "yuv420p10le" if codec == "av1" else "yuv420p" - stream.bit_rate = 0 - stream.options = {'crf': str(crf)} - if codec == "av1": - stream.options["preset"] = "6" + codec_map = {"vp9": "libvpx-vp9", "av1": "libsvtav1"} + stream = container.add_stream(codec_map[codec], rate=Fraction(round(fps * 1000), 1000)) + stream.width = images.shape[-2] + stream.height = images.shape[-3] + stream.pix_fmt = "yuv420p10le" if codec == "av1" else "yuv420p" + stream.bit_rate = 0 + stream.options = {'crf': str(crf)} + if codec == "av1": + stream.options["preset"] = "6" - for frame in images: - frame = av.VideoFrame.from_ndarray(torch.clamp(frame[..., :3] * 255, min=0, max=255).to(device=torch.device("cpu"), dtype=torch.uint8).numpy(), format="rgb24") - for packet in stream.encode(frame): - container.mux(packet) - container.mux(stream.encode()) - container.close() + for frame in images: + frame = av.VideoFrame.from_ndarray(torch.clamp(frame[..., :3] * 255, min=0, max=255).to(device=torch.device("cpu"), dtype=torch.uint8).numpy(), format="rgb24") + for packet in stream.encode(frame): + container.mux(packet) + container.mux(stream.encode()) + finally: + container.close() return io.NodeOutput(ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))