309 lines
12 KiB
Python
309 lines
12 KiB
Python
from __future__ import annotations
|
|
from typing_extensions import override
|
|
from typing import Optional
|
|
from fractions import Fraction
|
|
from enum import Enum
|
|
from PIL import Image
|
|
import torch
|
|
import numpy as np
|
|
import io
|
|
import json
|
|
import math
|
|
import av
|
|
import os
|
|
import random
|
|
import folder_paths
|
|
from comfy_api.latest._input import AudioInput, VideoInput
|
|
from comfy_api.latest import ComfyExtension, io, ui
|
|
from comfy_api.util import VideoComponents
|
|
|
|
class SaveVideoRGBA(io.ComfyNode):
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="SaveVideoRGBA",
|
|
display_name="Save Video (RGBA)",
|
|
category="image/animation",
|
|
inputs=[
|
|
io.Image.Input("images"),
|
|
io.Float.Input("fps", default=24.0, min=1.0, max=120.0, step=1.0),
|
|
io.String.Input("filename_prefix", default="video/ComfyUI", tooltip="The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."),
|
|
io.Boolean.Input("only_preview", default=False),
|
|
io.Audio.Input("audio", optional=True),
|
|
],
|
|
outputs=[
|
|
# io.Video.Output("VIDEO")
|
|
],
|
|
hidden=[io.Hidden.unique_id],
|
|
is_output_node=True
|
|
)
|
|
|
|
@classmethod
|
|
def execute(self, images, fps, filename_prefix, only_preview, audio=None, **kwargs):
|
|
|
|
B, H, W, C = images.shape
|
|
has_alpha = C == 4
|
|
|
|
divisible_by = 2
|
|
if W % divisible_by != 0 or H % divisible_by != 0:
|
|
new_width = W - (W % divisible_by)
|
|
new_height = H - (H % divisible_by)
|
|
|
|
print(f'Resize video from {W}x{H} to {new_width}x{new_height}')
|
|
|
|
resized_images = []
|
|
for i in range(B):
|
|
img = Image.fromarray((images[i].cpu().numpy() * 255).astype(np.uint8))
|
|
resized_img = img.resize((new_width, new_height), Image.LANCZOS)
|
|
resized_tensor = torch.from_numpy(np.array(resized_img).astype(np.float32) / 255.0)
|
|
resized_images.append(resized_tensor.unsqueeze(0))
|
|
|
|
images = torch.cat(resized_images, dim=0)
|
|
try:
|
|
from comfy_api.latest import InputImpl, Types
|
|
video = InputImpl.VideoFromComponents(Types.VideoComponents(images=images, audio=audio, frame_rate=Fraction(fps)))
|
|
except Exception as e:
|
|
video = RGBAVideoFromComponents(
|
|
VideoComponents(
|
|
images=images,
|
|
audio=audio,
|
|
frame_rate=Fraction(fps),
|
|
)
|
|
)
|
|
|
|
width, height = video.get_dimensions()
|
|
results = list()
|
|
|
|
# 预览
|
|
if only_preview or has_alpha:
|
|
output_dir = folder_paths.get_temp_directory()
|
|
prefix_append = "ComfyUI_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5))
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
|
|
prefix_append,
|
|
output_dir,
|
|
width,
|
|
height
|
|
)
|
|
format = 'webm' if has_alpha else 'auto'
|
|
file = f"{filename}_{counter:05}_.{RGBAVideoContainer.get_extension(format)}"
|
|
video.save_to(
|
|
path=os.path.join(full_output_folder, file),
|
|
format=format,
|
|
codec='auto',
|
|
metadata=None,
|
|
)
|
|
|
|
results.append(ui.SavedResult(file, subfolder, io.FolderType.temp))
|
|
counter += 1
|
|
|
|
# 保存
|
|
if not only_preview:
|
|
output_dir = folder_paths.get_output_directory()
|
|
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
|
|
filename_prefix,
|
|
output_dir,
|
|
width,
|
|
height
|
|
)
|
|
format = 'mov' if has_alpha else 'auto'
|
|
codec = 'prores_ks' if has_alpha else 'auto'
|
|
file = f"{filename}_{counter:05}_.{RGBAVideoContainer.get_extension(format)}"
|
|
video.save_to(
|
|
path=os.path.join(full_output_folder, file),
|
|
format=format,
|
|
codec=codec,
|
|
metadata=None,
|
|
)
|
|
|
|
if not has_alpha:
|
|
results.append(ui.SavedResult(file, subfolder, io.FolderType.output))
|
|
counter += 1
|
|
|
|
return io.NodeOutput(ui=ui.PreviewVideo(results))
|
|
|
|
class NodeExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
SaveVideoRGBA
|
|
]
|
|
|
|
async def comfy_entrypoint() -> NodeExtension:
|
|
return NodeExtension()
|
|
|
|
|
|
|
|
class RGBAVideoCodec(str, Enum):
|
|
AUTO = "auto"
|
|
H264 = "h264" # H.264 codec, no alpha channel
|
|
VP9 = "libvpx-vp9" # VP9 codec, supports alpha channel
|
|
PRORES = "prores_ks" # ProRes 4444, supports alpha channel
|
|
|
|
@classmethod
|
|
def as_input(cls) -> list[str]:
|
|
"""
|
|
Returns a list of codec names that can be used as node input.
|
|
"""
|
|
return [member.value for member in cls]
|
|
|
|
class RGBAVideoContainer(str, Enum):
|
|
AUTO = "auto"
|
|
MP4 = "mp4" # MP4 container
|
|
WEBM = "webm" # WebM container, supports VP9 with alpha channel
|
|
MOV = "mov" # QuickTime MOV, supports ProRes with alpha channel
|
|
|
|
@classmethod
|
|
def as_input(cls) -> list[str]:
|
|
"""
|
|
Returns a list of codec names that can be used as node input.
|
|
"""
|
|
return [member.value for member in cls]
|
|
|
|
@classmethod
|
|
def get_extension(cls, value) -> str:
|
|
"""
|
|
Returns the file extension for the container.
|
|
"""
|
|
if isinstance(value, str):
|
|
value = cls(value)
|
|
if value == RGBAVideoContainer.AUTO:
|
|
return "mp4"
|
|
elif value == RGBAVideoContainer.MP4:
|
|
return "mp4"
|
|
elif value == RGBAVideoContainer.WEBM:
|
|
return "webm"
|
|
elif value == RGBAVideoContainer.MOV:
|
|
return "mov"
|
|
return ""
|
|
|
|
class RGBAVideoFromComponents(VideoInput):
|
|
|
|
def __init__(self, components: VideoComponents):
|
|
self.__components = components
|
|
|
|
def get_components(self) -> VideoComponents:
|
|
return VideoComponents(
|
|
images=self.__components.images,
|
|
audio=self.__components.audio,
|
|
frame_rate=self.__components.frame_rate
|
|
)
|
|
|
|
def save_to(
|
|
self,
|
|
path: str,
|
|
format: RGBAVideoContainer = RGBAVideoContainer.AUTO,
|
|
codec: RGBAVideoCodec = RGBAVideoCodec.AUTO,
|
|
metadata: Optional[dict] = None
|
|
):
|
|
# Check if images have alpha channel (4 channels)
|
|
has_alpha = self.__components.images.shape[-1] == 4 if len(self.__components.images.shape) == 4 else False
|
|
|
|
# Determine format and codec based on alpha channel
|
|
if has_alpha:
|
|
# For alpha channel support, use webm container with vp9 codec
|
|
if format == RGBAVideoContainer.AUTO:
|
|
format = RGBAVideoContainer.WEBM.value
|
|
if codec == RGBAVideoCodec.AUTO:
|
|
codec = RGBAVideoCodec.VP9.value
|
|
# Ensure format is string for comparison
|
|
format_str = format.value if isinstance(format, RGBAVideoContainer) else format
|
|
codec_str = codec.value if isinstance(codec, RGBAVideoCodec) else codec
|
|
if format_str not in ['webm', 'mov']:
|
|
raise ValueError("Only WEBM and MOV formats support alpha channel")
|
|
else:
|
|
# Ensure format and codec are strings for comparison
|
|
format_str = format.value if isinstance(format, RGBAVideoContainer) else format
|
|
codec_str = codec.value if isinstance(codec, RGBAVideoCodec) else codec
|
|
if format != RGBAVideoContainer.AUTO and format_str not in ['mp4', 'webm', 'mov']:
|
|
raise ValueError("Supported formats: MP4, WEBM, MOV")
|
|
if codec != RGBAVideoCodec.AUTO and codec_str not in ['h264', 'libvpx-vp9', 'prores_ks']:
|
|
raise ValueError("Supported codecs: H264, VP9, ProRes")
|
|
|
|
# Set default codec if AUTO
|
|
if codec == RGBAVideoCodec.AUTO:
|
|
codec = RGBAVideoCodec.H264.value
|
|
codec_str = codec
|
|
|
|
# Ensure codec_str is defined
|
|
if 'codec_str' not in locals():
|
|
codec_str = codec.value if isinstance(codec, RGBAVideoCodec) else codec
|
|
|
|
# Prepare options based on format
|
|
options = {}
|
|
if format_str in ['mp4', 'mov']:
|
|
options['movflags'] = 'use_metadata_tags'
|
|
|
|
# Determine the format string for av.open
|
|
output_format = None if format == RGBAVideoContainer.AUTO else format_str
|
|
|
|
with av.open(path, mode='w', format=output_format, options=options) as output:
|
|
# Add metadata before writing any streams
|
|
if metadata is not None:
|
|
for key, value in metadata.items():
|
|
output.metadata[key] = json.dumps(value)
|
|
|
|
frame_rate = Fraction(round(self.__components.frame_rate * 1000), 1000)
|
|
|
|
# Create a video stream
|
|
video_stream = output.add_stream(codec_str, rate=frame_rate)
|
|
video_stream.width = self.__components.images.shape[2]
|
|
video_stream.height = self.__components.images.shape[1]
|
|
|
|
# Set pixel format based on codec and alpha channel
|
|
if has_alpha:
|
|
if codec_str in ['libvpx-vp9']:
|
|
video_stream.pix_fmt = 'yuva420p' # VP9 with alpha
|
|
elif codec_str in ['prores_ks']:
|
|
video_stream.pix_fmt = 'yuva444p10le' # ProRes 4444
|
|
else:
|
|
video_stream.pix_fmt = 'yuva420p' # Default alpha format
|
|
else:
|
|
if codec_str in ['h264']:
|
|
video_stream.pix_fmt = 'yuv420p'
|
|
elif codec_str in ['libvpx-vp9']:
|
|
video_stream.pix_fmt = 'yuv420p'
|
|
else:
|
|
video_stream.pix_fmt = 'yuv420p'
|
|
|
|
# Create an audio stream
|
|
audio_sample_rate = 1
|
|
audio_stream: Optional[av.AudioStream] = None
|
|
if self.__components.audio:
|
|
audio_sample_rate = int(self.__components.audio['sample_rate'])
|
|
audio_codec = 'libopus' if format_str == 'webm' else 'aac'
|
|
audio_stream = output.add_stream(audio_codec, rate=audio_sample_rate)
|
|
|
|
# Encode video
|
|
for i, frame in enumerate(self.__components.images):
|
|
img = (frame * 255).clamp(0, 255).byte().cpu().numpy()
|
|
|
|
if has_alpha:
|
|
# Create frame with alpha channel (H, W, 4)
|
|
frame = av.VideoFrame.from_ndarray(img, format='rgba')
|
|
frame = frame.reformat(format=video_stream.pix_fmt)
|
|
else:
|
|
# Create frame without alpha channel (H, W, 3)
|
|
frame = av.VideoFrame.from_ndarray(img, format='rgb24')
|
|
frame = frame.reformat(format=video_stream.pix_fmt)
|
|
|
|
packet = video_stream.encode(frame)
|
|
output.mux(packet)
|
|
|
|
# Flush video
|
|
packet = video_stream.encode(None)
|
|
output.mux(packet)
|
|
|
|
if audio_stream and self.__components.audio:
|
|
waveform = self.__components.audio['waveform']
|
|
waveform = waveform[:, :,
|
|
:math.ceil((audio_sample_rate / frame_rate) * self.__components.images.shape[0])]
|
|
frame = av.AudioFrame.from_ndarray(waveform.movedim(2, 1).reshape(1, -1).float().numpy(), format='flt',
|
|
layout='mono' if waveform.shape[1] == 1 else 'stereo')
|
|
frame.sample_rate = audio_sample_rate
|
|
frame.pts = 0
|
|
output.mux(audio_stream.encode(frame))
|
|
|
|
# Flush encoder
|
|
output.mux(audio_stream.encode(None)) |