refactor(cryptio): 优化视频保存逻辑，统一处理重封装和重新编码

src/cryptio/save_video.py

@@ -3,11 +3,139 @@ import random
 from io import BytesIO
 import folder_paths  # pyright: ignore[reportMissingImports]
 from .utils import _key_manager
+import av
+import torch
+import numpy as np
+import math
+import json
+from fractions import Fraction
+from .utils import _key_manager
 from .utils.crypto_utils import encrypt_data_hybrid
 from comfy_api.input import VideoInput
 from comfy_api.latest import io, ui
 
 
+
+def _save_to_bytesio_impl(video, buffer, metadata=None):
+    """
+    Unified method to save video (file or components) to a BytesIO buffer.
+    Handles remuxing for files (if compatible) and re-encoding for components.
+    """
+    # 1. Try Remuxing (if input is a file)
+    if hasattr(video, "get_stream_source"):
+        try:
+            source = video.get_stream_source()
+            if isinstance(source, BytesIO):
+                source.seek(0)
+            
+            with av.open(source, mode='r') as input_container:
+                # We strictly want MP4 for encryption node compatibility
+                # writing to buffer requires explicit format
+                with av.open(buffer, mode='w', format='mp4', options={'movflags': 'use_metadata_tags'}) as output_container:
+                    # Merge metadata
+                    if metadata is None:
+                        metadata = {}
+                    
+                    # Copy input metadata first
+                    for key, value in input_container.metadata.items():
+                        if key not in metadata:
+                            output_container.metadata[key] = value
+                    
+                    # Apply new metadata
+                    for key, value in metadata.items():
+                        output_container.metadata[key] = json.dumps(value) if not isinstance(value, str) else value
+
+                    # Map streams
+                    stream_map = {}
+                    for stream in input_container.streams:
+                        # Copy video and audio streams
+                        if isinstance(stream, (av.VideoStream, av.AudioStream)):
+                            out_stream = output_container.add_stream_from_template(template=stream, opaque=True)
+                            stream_map[stream] = out_stream
+                    
+                    if not stream_map:
+                         raise ValueError("No compatible streams found for remuxing")
+
+                    # Copy packets
+                    for packet in input_container.demux():
+                        if packet.stream in stream_map and packet.dts is not None:
+                            packet.stream = stream_map[packet.stream]
+                            output_container.mux(packet)
+            
+            return # Success
+        except Exception as e:
+            # Fallback to re-encoding
+            # print(f"CryptIO: Remux failed ({e}), falling back to re-encoding.")
+            buffer.seek(0)
+            buffer.truncate()
+            pass
+
+    # 2. Re-encoding (VideoFromComponents or fallback)
+    components = video.get_components()
+    
+    with av.open(buffer, mode='w', format='mp4', options={'movflags': 'use_metadata_tags'}) as output:
+        if metadata:
+            for key, value in metadata.items():
+                output.metadata[key] = json.dumps(value) if not isinstance(value, str) else value
+
+        # Video Stream
+        frame_rate = components.frame_rate
+        if isinstance(frame_rate, (int, float)):
+             frame_rate = Fraction(round(frame_rate * 1000), 1000)
+             
+        if components.images is None or len(components.images) == 0:
+             raise ValueError("No video frames to save")
+
+        height, width = components.images.shape[1], components.images.shape[2]
+        
+        stream = output.add_stream('h264', rate=frame_rate)
+        stream.width = width
+        stream.height = height
+        stream.pix_fmt = 'yuv420p'
+
+        # Audio Stream
+        audio_stream = None
+        audio_sample_rate = 44100
+        if components.audio:
+            audio_sample_rate = int(components.audio.get('sample_rate', 44100))
+            audio_stream = output.add_stream('aac', rate=audio_sample_rate)
+
+        # Encode Video
+        for i, img_tensor in enumerate(components.images):
+            img_np = (img_tensor * 255).clamp(0, 255).byte().cpu().numpy()
+            frame = av.VideoFrame.from_ndarray(img_np, format='rgb24')
+            frame = frame.reformat(format='yuv420p')
+            for packet in stream.encode(frame):
+                output.mux(packet)
+        
+        # Flush Video
+        for packet in stream.encode(None):
+            output.mux(packet)
+            
+        # Encode Audio
+        if audio_stream and components.audio:
+            waveform = components.audio['waveform']
+            required_samples = math.ceil((audio_sample_rate / frame_rate) * len(components.images))
+            if waveform.shape[-1] > required_samples:
+                waveform = waveform[..., :required_samples]
+            
+            flat_audio = waveform.movedim(2, 1).reshape(1, -1).float().cpu().numpy()
+            
+            frame = av.AudioFrame.from_ndarray(
+                flat_audio, 
+                format='flt', 
+                layout='mono' if waveform.shape[1] == 1 else 'stereo'
+            )
+            frame.sample_rate = audio_sample_rate
+            frame.pts = 0
+            
+            for packet in audio_stream.encode(frame):
+                output.mux(packet)
+            
+            for packet in audio_stream.encode(None):
+                output.mux(packet)
+
+
 class SaveVideoCryptIO(io.ComfyNode):
     """
     Encrypt and save video files using client public key
@@ -55,7 +183,9 @@ class SaveVideoCryptIO(io.ComfyNode):
 
             # Save video to in-memory buffer to avoid writing plaintext to disk
             buffer = BytesIO()
-            video.save_to(buffer, format="mp4", codec="auto", metadata=metadata)
+            # Save video to in-memory buffer using unified helper
+            buffer = BytesIO()
+            _save_to_bytesio_impl(video, buffer, metadata=metadata)
             video_data = buffer.getvalue()
             buffer.close()
 
@@ -130,9 +260,10 @@ class PreviewVideoCryptIO(io.ComfyNode):
                 if cls.hidden.prompt is not None:
                     metadata["prompt"] = cls.hidden.prompt
 
-            # Save video to in-memory buffer to avoid writing plaintext to disk
+            # Save video to in-memory buffer using unified helper
             buffer = BytesIO()
-            video.save_to(buffer, format="mp4", codec="auto", metadata=metadata)
+            # Preview forces mp4/h264 in save_to usually, but our helper does that by default for buffer
+            _save_to_bytesio_impl(video, buffer, metadata=metadata)
             video_data = buffer.getvalue()
             buffer.close()