Preserving hardware memory during hw decoding, exporting/importing via dlpack (#2155)

Add DLPack support for zero-copy GPU frame exchange between PyAV and frameworks like PyTorch/CuPy, avoiding GPU->CPU->GPU round-trips when using cuvid decode + GPU processing + nvenc encode.

- Implement __dlpack__ and __dlpack_device__ on VideoPlane
- Add VideoFrame.from_dlpack for importing DLPack tensors (nv12/p010le/p016le)
- Add CudaContext for managing CUDA device/frames contexts
- Add HWAccel.is_hw_owned to preserve frames on GPU memory
- Switch frame buffer management from av_image_alloc to av_frame_get_buffer
- Handle hw_frames_ctx in planes, reformatter, and to_ndarray paths

Author: caffeinism

av/codec/hwaccel.pxd

@@ -16,6 +16,8 @@ cdef class HWAccel:
     cdef readonly Codec codec
     cdef readonly HWConfig config
     cdef lib.AVBufferRef *ptr
+    cdef readonly int device_id
+    cdef readonly bint is_hw_owned
     cdef public bint allow_software_fallback
     cdef public dict options
     cdef public int flags

av/codec/hwaccel.py

@@ -110,6 +110,7 @@ def __init__(
         allow_software_fallback=True,
         options=None,
         flags=None,
+        is_hw_owned=False,
     ):
         if isinstance(device_type, HWDeviceType):
             self._device_type = device_type
@@ -120,9 +121,17 @@ def __init__(
         else:
             raise ValueError("Unknown type for device_type")
 
-        self._device = device
+        self.is_hw_owned = is_hw_owned
+        self.device_id = 0
+        if self._device_type == HWDeviceType.cuda and device:
+            self.device_id = int(device)
+
+        self._device = None if device is None else f"{device}"
         self.allow_software_fallback = allow_software_fallback
+
         self.options = {} if not options else dict(options)
+        if self._device_type == HWDeviceType.cuda and self.is_hw_owned:
+            self.options.setdefault("primary_ctx", "1")
         self.flags = 0 if not flags else flags
         self.ptr = cython.NULL
         self.config = None
@@ -135,7 +144,7 @@ def _initialize_hw_context(self, codec: Codec):
             if self._device_type and config.device_type != self._device_type:
                 continue
             break
-        else:
+        else:  # nobreak
             raise NotImplementedError(f"No supported hardware config for {codec}")
 
         self.config = config
@@ -155,7 +164,7 @@ def _initialize_hw_context(self, codec: Codec):
             )
         )
 
-    def create(self, codec: Codec):
+    def create(self, codec: Codec) -> HWAccel:
         """Create a new hardware accelerator context with the given codec"""
         if self.ptr:
             raise RuntimeError("Hardware context already initialized")
@@ -165,6 +174,7 @@ def create(self, codec: Codec):
             device=self._device,
             allow_software_fallback=self.allow_software_fallback,
             options=self.options,
+            is_hw_owned=self.is_hw_owned,
         )
         ret._initialize_hw_context(codec)
         return ret

av/codec/hwaccel.pyi

@@ -37,13 +37,20 @@ class HWConfig:
     def is_supported(self) -> bool: ...
 
 class HWAccel:
+    options: dict[str, object]
+
+    @property
+    def is_hw_owned(self) -> bool: ...
+    @property
+    def device_id(self) -> int: ...
     def __init__(
         self,
         device_type: str | HWDeviceType,
-        device: str | None = None,
+        device: str | int | None = None,
         allow_software_fallback: bool = False,
         options: dict[str, object] | None = None,
         flags: int | None = None,
+        is_hw_owned: bool = False,
     ) -> None: ...
     def create(self, codec: Codec) -> HWAccel: ...
 

av/video/codeccontext.py

@@ -127,6 +127,10 @@ def _transfer_hwframe(self, frame: Frame):
             # need to transfer.
             return frame
 
+        if self.hwaccel_ctx.is_hw_owned:
+            cython.cast(VideoFrame, frame)._device_id = self.hwaccel_ctx.device_id
+            return frame
+
         frame_sw: Frame = self._alloc_next_frame()
         err_check(lib.av_hwframe_transfer_data(frame_sw.ptr, frame.ptr, 0))
         # TODO: Is there anything else to transfer?

av/video/frame.pxd

@@ -6,15 +6,21 @@ from av.video.format cimport VideoFormat
 from av.video.reformatter cimport VideoReformatter
 
 
-cdef class VideoFrame(Frame):
-    # This is the buffer that is used to back everything in the AVFrame.
-    # We don't ever actually access it directly.
-    cdef uint8_t *_buffer
-    cdef object _np_buffer
+cdef class CudaContext:
+    cdef readonly int device_id
+    cdef readonly bint primary_ctx
+    cdef lib.AVBufferRef* _device_ref
+    cdef dict _frames_cache
+    cdef lib.AVBufferRef* _get_device_ref(self)
+    cdef public lib.AVBufferRef* get_frames_ctx(
+        self, lib.AVPixelFormat sw_fmt, int width, int height
+    )
 
+cdef class VideoFrame(Frame):
+    cdef CudaContext _cuda_ctx
     cdef VideoReformatter reformatter
     cdef readonly VideoFormat format
-
+    cdef readonly int _device_id
     cdef _init(self, lib.AVPixelFormat format, unsigned int width, unsigned int height)
     cdef _init_user_attributes(self)
     cpdef save(self, object filepath)