From be4b655f016ddc826c17de930e2301ab30704975 Mon Sep 17 00:00:00 2001
From: bilt <vanderbilt@gmail.com>
Date: Thu, 31 Jul 2025 18:18:58 +0800
Subject: [PATCH] SaveVideo

---
 comfy_api/input/video_types.py      | 15 +++++++++++++
 comfy_api/input_impl/video_types.py | 21 +++++++++++++++++-
 comfy_extras/nodes_video.py         | 33 +++++++++--------------------
 3 files changed, 45 insertions(+), 24 deletions(-)

diff --git a/comfy_api/input/video_types.py b/comfy_api/input/video_types.py
index dc22d34ff..dd3fd5c62 100644
--- a/comfy_api/input/video_types.py
+++ b/comfy_api/input/video_types.py
@@ -18,6 +18,21 @@ class VideoInput(ABC):
         """
         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
     def save_to(
         self,
diff --git a/comfy_api/input_impl/video_types.py b/comfy_api/input_impl/video_types.py
index 197f6558c..ebcbc8c8e 100644
--- a/comfy_api/input_impl/video_types.py
+++ b/comfy_api/input_impl/video_types.py
@@ -159,6 +159,16 @@ class VideoFromFile(VideoInput):
             return self.get_components_internal(container)
         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(
         self,
         path: str | io.BytesIO,
@@ -218,6 +228,15 @@ class VideoFromFile(VideoInput):
                         output_container.mux(packet)
 
 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.
     """
@@ -234,7 +253,7 @@ class VideoFromComponents(VideoInput):
 
     def save_to(
         self,
-        path: str,
+        path: str | io.BytesIO,
         format: VideoContainer = VideoContainer.AUTO,
         codec: VideoCodec = VideoCodec.AUTO,
         metadata: Optional[dict] = None
diff --git a/comfy_extras/nodes_video.py b/comfy_extras/nodes_video.py
index b608aacff..602f3dae8 100644
--- a/comfy_extras/nodes_video.py
+++ b/comfy_extras/nodes_video.py
@@ -158,35 +158,22 @@ class SaveVideo(ComfyNodeABC):
         
         # 生成文件路径
         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}")
-        
-        # 先保存到临时文件
-        video.save_to(
-            temp_file_path,
-            format=format,
-            codec=codec,
-            metadata=saved_metadata
-        )
-        
-        # 如果启用加密，对文件进行XOR加密
+
         if encrypt:
-            # 读取文件内容
-            with open(temp_file_path, 'rb') as f:
-                video_data = f.read()
-            
-            # 删除临时文件
-            os.remove(temp_file_path)
-            
-            # 加密数据
+            # 直接在内存中对视频数据进行加密
+            video_data = video.get_video_data(format=format, codec=codec, metadata=saved_metadata)
             encrypted_data = self.simple_xor_encrypt(video_data, encryption_key)
-            
-            # 保存加密后的文件
             with open(final_file_path, 'wb') as f:
                 f.write(encrypted_data)
         else:
-            # 不加密，直接重命名
-            os.rename(temp_file_path, final_file_path)
+            # 不加密，直接保存
+            video.save_to(
+                final_file_path,
+                format=format,
+                codec=codec,
+                metadata=saved_metadata
+            )
 
         # 上传到Hugging Face
         self.upload(args.hf_token, args.hf_dataset_name, final_file_path)