mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-08-25 18:05:45 +08:00
SaveVideo
This commit is contained in:
parent
ffa1b09ae7
commit
be4b655f01
@ -18,6 +18,21 @@ class VideoInput(ABC):
|
|||||||
"""
|
"""
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def get_video_data(
|
||||||
|
self,
|
||||||
|
format: VideoContainer = VideoContainer.AUTO,
|
||||||
|
codec: VideoCodec = VideoCodec.AUTO,
|
||||||
|
metadata: Optional[dict] = None
|
||||||
|
) -> bytes:
|
||||||
|
"""
|
||||||
|
Abstract method to get the video data as bytes.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Video data as bytes
|
||||||
|
"""
|
||||||
|
pass
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def save_to(
|
def save_to(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@ -159,6 +159,16 @@ class VideoFromFile(VideoInput):
|
|||||||
return self.get_components_internal(container)
|
return self.get_components_internal(container)
|
||||||
raise ValueError(f"No video stream found in file '{self.__file}'")
|
raise ValueError(f"No video stream found in file '{self.__file}'")
|
||||||
|
|
||||||
|
def get_video_data(
|
||||||
|
self,
|
||||||
|
format: VideoContainer = VideoContainer.AUTO,
|
||||||
|
codec: VideoCodec = VideoCodec.AUTO,
|
||||||
|
metadata: Optional[dict] = None
|
||||||
|
) -> bytes:
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
self.save_to(buffer, format, codec, metadata)
|
||||||
|
return buffer.getvalue()
|
||||||
|
|
||||||
def save_to(
|
def save_to(
|
||||||
self,
|
self,
|
||||||
path: str | io.BytesIO,
|
path: str | io.BytesIO,
|
||||||
@ -218,6 +228,15 @@ class VideoFromFile(VideoInput):
|
|||||||
output_container.mux(packet)
|
output_container.mux(packet)
|
||||||
|
|
||||||
class VideoFromComponents(VideoInput):
|
class VideoFromComponents(VideoInput):
|
||||||
|
def get_video_data(
|
||||||
|
self,
|
||||||
|
format: VideoContainer = VideoContainer.AUTO,
|
||||||
|
codec: VideoCodec = VideoCodec.AUTO,
|
||||||
|
metadata: Optional[dict] = None
|
||||||
|
) -> bytes:
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
self.save_to(buffer, format, codec, metadata)
|
||||||
|
return buffer.getvalue()
|
||||||
"""
|
"""
|
||||||
Class representing video input from tensors.
|
Class representing video input from tensors.
|
||||||
"""
|
"""
|
||||||
@ -234,7 +253,7 @@ class VideoFromComponents(VideoInput):
|
|||||||
|
|
||||||
def save_to(
|
def save_to(
|
||||||
self,
|
self,
|
||||||
path: str,
|
path: str | io.BytesIO,
|
||||||
format: VideoContainer = VideoContainer.AUTO,
|
format: VideoContainer = VideoContainer.AUTO,
|
||||||
codec: VideoCodec = VideoCodec.AUTO,
|
codec: VideoCodec = VideoCodec.AUTO,
|
||||||
metadata: Optional[dict] = None
|
metadata: Optional[dict] = None
|
||||||
|
|||||||
@ -158,35 +158,22 @@ class SaveVideo(ComfyNodeABC):
|
|||||||
|
|
||||||
# 生成文件路径
|
# 生成文件路径
|
||||||
file_extension = VideoContainer.get_extension(format)
|
file_extension = VideoContainer.get_extension(format)
|
||||||
temp_file_path = os.path.join(full_output_folder, f"temp_{filename}_{counter:05}_.{file_extension}")
|
|
||||||
final_file_path = os.path.join(full_output_folder, f"{filename}_{counter:05}_.{file_extension}")
|
final_file_path = os.path.join(full_output_folder, f"{filename}_{counter:05}_.{file_extension}")
|
||||||
|
|
||||||
# 先保存到临时文件
|
|
||||||
video.save_to(
|
|
||||||
temp_file_path,
|
|
||||||
format=format,
|
|
||||||
codec=codec,
|
|
||||||
metadata=saved_metadata
|
|
||||||
)
|
|
||||||
|
|
||||||
# 如果启用加密,对文件进行XOR加密
|
|
||||||
if encrypt:
|
if encrypt:
|
||||||
# 读取文件内容
|
# 直接在内存中对视频数据进行加密
|
||||||
with open(temp_file_path, 'rb') as f:
|
video_data = video.get_video_data(format=format, codec=codec, metadata=saved_metadata)
|
||||||
video_data = f.read()
|
|
||||||
|
|
||||||
# 删除临时文件
|
|
||||||
os.remove(temp_file_path)
|
|
||||||
|
|
||||||
# 加密数据
|
|
||||||
encrypted_data = self.simple_xor_encrypt(video_data, encryption_key)
|
encrypted_data = self.simple_xor_encrypt(video_data, encryption_key)
|
||||||
|
|
||||||
# 保存加密后的文件
|
|
||||||
with open(final_file_path, 'wb') as f:
|
with open(final_file_path, 'wb') as f:
|
||||||
f.write(encrypted_data)
|
f.write(encrypted_data)
|
||||||
else:
|
else:
|
||||||
# 不加密,直接重命名
|
# 不加密,直接保存
|
||||||
os.rename(temp_file_path, final_file_path)
|
video.save_to(
|
||||||
|
final_file_path,
|
||||||
|
format=format,
|
||||||
|
codec=codec,
|
||||||
|
metadata=saved_metadata
|
||||||
|
)
|
||||||
|
|
||||||
# 上传到Hugging Face
|
# 上传到Hugging Face
|
||||||
self.upload(args.hf_token, args.hf_dataset_name, final_file_path)
|
self.upload(args.hf_token, args.hf_dataset_name, final_file_path)
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user