fix VideoFromFile stream source to _ReentrantBytesIO for parallel async use (with pickle support)

Author: bigcat88

comfy_api/latest/_input_impl/video_types.py

@@ -2,8 +2,8 @@
 from av.container import InputContainer
 from av.subtitles.stream import SubtitleStream
 from fractions import Fraction
-from typing import Optional, IO, Iterator
-from contextlib import contextmanager, suppress
+from typing import Optional, Iterator
+from contextlib import contextmanager
 from .._input import AudioInput, VideoInput
 import av
 import io
@@ -22,8 +22,8 @@ def __init__(self, data: bytes):
         if data is None:
             raise TypeError("data must be bytes, not None")
         self._data = data
-        self._view = memoryview(data)
         self._pos = 0
+        self._len = len(data)
 
     def getvalue(self) -> bytes:
         if self.closed:
@@ -33,7 +33,7 @@ def getvalue(self) -> bytes:
     def getbuffer(self) -> memoryview:
         if self.closed:
             raise ValueError("I/O operation on closed file.")
-        return memoryview(self._data)  # return a NEW view; external .release() won't break our internal _view.
+        return memoryview(self._data)
 
     def readable(self) -> bool:
         return True
@@ -55,7 +55,7 @@ def seek(self, offset: int, whence: int = io.SEEK_SET) -> int:
         elif whence == io.SEEK_CUR:
             new_pos = self._pos + offset
         elif whence == io.SEEK_END:
-            new_pos = len(self._view) + offset
+            new_pos = self._len + offset
         else:
             raise ValueError(f"Invalid whence: {whence}")
         if new_pos < 0:
@@ -70,10 +70,10 @@ def readinto(self, b) -> int:
         if mv.readonly:
             raise TypeError("readinto() argument must be writable")
         mv = mv.cast("B")
-        if self._pos >= len(self._view):
+        if self._pos >= self._len:
             return 0
-        n = min(len(mv), len(self._view) - self._pos)
-        mv[:n] = self._view[self._pos:self._pos + n]
+        n = min(len(mv), self._len - self._pos)
+        mv[:n] = self._data[self._pos:self._pos + n]
         self._pos += n
         return n
 
@@ -84,10 +84,10 @@ def read(self, size: int = -1) -> bytes:
         if self.closed:
             raise ValueError("I/O operation on closed file.")
         if size is None or size < 0:
-            size = len(self._view) - self._pos
-        if self._pos >= len(self._view):
+            size = self._len - self._pos
+        if self._pos >= self._len:
             return b""
-        end = min(self._pos + size, len(self._view))
+        end = min(self._pos + size, self._len)
         out = self._data[self._pos:end]
         self._pos = end
         return out
@@ -98,9 +98,9 @@ def read1(self, size: int = -1) -> bytes:
     def readline(self, size: int = -1) -> bytes:
         if self.closed:
             raise ValueError("I/O operation on closed file.")
-        if self._pos >= len(self._view):
+        if self._pos >= self._len:
             return b""
-        end_limit = len(self._view) if size is None or size < 0 else min(len(self._view), self._pos + size)
+        end_limit = self._len if size is None or size < 0 else min(self._len, self._pos + size)
         nl = self._data.find(b"\n", self._pos, end_limit)
         end = (nl + 1) if nl != -1 else end_limit
         out = self._data[self._pos:end]
@@ -131,11 +131,6 @@ def writelines(self, lines) -> None:
     def truncate(self, size: int | None = None) -> int:
         raise io.UnsupportedOperation("not writable")
 
-    def close(self) -> None:
-        with suppress(Exception):
-            self._view.release()
-        super().close()
-
 
 def container_to_output_format(container_format: str | None) -> str | None:
     """
@@ -181,15 +176,15 @@ class VideoFromFile(VideoInput):
     Class representing video input from a file.
     """
 
+    __data: str | bytes
+
     def __init__(self, file: str | io.BytesIO | bytes | bytearray | memoryview):
         """
         Initialize the VideoFromFile object based off of either a path on disk or a BytesIO object
         containing the file contents.
         """
-        self.__path: Optional[str] = None
-        self.__data: Optional[bytes] = None
         if isinstance(file, str):
-            self.__path = file
+            self.__data = file
         elif isinstance(file, io.BytesIO):
             # Snapshot to immutable bytes once to ensure re-entrant, parallel-safe readers.
             self.__data = file.getbuffer().tobytes()
@@ -203,12 +198,9 @@ def get_stream_source(self) -> str | io.BytesIO:
         Return the underlying file source for efficient streaming.
         This avoids unnecessary memory copies when the source is already a file path.
         """
-        if self.__path is not None:
-            return self.__path
-        data = self.__data
-        if data is None:
-            raise RuntimeError("VideoFromFile: missing in-memory bytes (__data is None)")
-        return _ReentrantBytesIO(data)
+        if isinstance(self.__data, str):
+            return self.__data
+        return _ReentrantBytesIO(self.__data)
 
     def get_dimensions(self) -> tuple[int, int]:
         """
@@ -430,20 +422,11 @@ def _get_first_video_stream(self, container: InputContainer):
         return video_stream
 
     def _source_label(self) -> str:
-        if self.__path is not None:
-            return self.__path
-        return f"<in-memory video: {len(self.__data)} bytes>"
+        return self.__data if isinstance(self.__data, str) else f"<in-memory video: {len(self.__data)} bytes>"
 
     @contextmanager
-    def _open_source(self) -> Iterator[str | IO[bytes]]:
-        """Internal helper to ensure file-like sources are closed after use."""
-        src = self.get_stream_source()
-        try:
-            yield src
-        finally:
-            if not isinstance(src, str):
-                with suppress(Exception):
-                    src.close()
+    def _open_source(self) -> Iterator[str | io.BytesIO]:
+        yield self.get_stream_source()
 
 
 class VideoFromComponents(VideoInput):