feat(cuda): support current CUDA contexts

av/video/frame.pxd

@@ -9,6 +9,7 @@ from av.video.reformatter cimport VideoReformatter
 cdef class CudaContext:
     cdef readonly int device_id
     cdef readonly bint primary_ctx
+    cdef readonly bint current_ctx
     cdef lib.AVBufferRef* _device_ref
     cdef dict _frames_cache
     cdef lib.AVBufferRef* _get_device_ref(self)

av/video/frame.py

@@ -22,11 +22,28 @@
 @cython.final
 @cython.cclass
 class CudaContext:
-    def __cinit__(self, device_id: cython.int = 0, primary_ctx: cython.bint = True):
+    """A reusable FFmpeg CUDA context for frames imported with DLPack.
+
+    :param int device_id: CUDA device ordinal.
+    :param bool primary_ctx: Retain the device's primary CUDA context.
+    :param bool current_ctx: Wrap the CUDA context current on the calling thread
+        when this object is first used. This is mutually exclusive with
+        ``primary_ctx``.
+    """
+
+    def __cinit__(
+        self,
+        device_id: cython.int = 0,
+        primary_ctx: cython.bint = True,
+        current_ctx: cython.bint = False,
+    ):
         self.device_id = device_id
         self.primary_ctx = primary_ctx
+        self.current_ctx = current_ctx
         self._device_ref = cython.NULL
         self._frames_cache = {}
+        if primary_ctx and current_ctx:
+            raise ValueError("primary_ctx and current_ctx are mutually exclusive")
 
     def __dealloc__(self):
         ref: cython.pointer[lib.AVBufferRef]
@@ -53,9 +70,10 @@ def _get_device_ref(self) -> cython.pointer[lib.AVBufferRef]:
         device_ref = cython.NULL
         device_bytes = f"{self.device_id}".encode()
         c_device: cython.p_char = device_bytes
-        options: Dictionary = Dictionary(
-            {"primary_ctx": "1" if self.primary_ctx else "0"}
-        )
+        options_dict = {"primary_ctx": "1" if self.primary_ctx else "0"}
+        if self.current_ctx:
+            options_dict["current_ctx"] = "1"
+        options: Dictionary = Dictionary(options_dict)
         err_check(
             lib.av_hwdevice_ctx_create(
                 cython.address(device_ref),
@@ -1471,6 +1489,7 @@ def from_dlpack(
         device_id: int | None = None,
         primary_ctx: bool = True,
         cuda_context=None,
+        current_ctx: bool = False,
     ):
         if not isinstance(planes, (tuple, list)):
             planes = (planes,)
@@ -1612,7 +1631,11 @@ def from_dlpack(
                 ctx: CudaContext
                 frames_ref: cython.pointer[lib.AVBufferRef]
                 if cuda_context is None:
-                    ctx = CudaContext(device_id=device_id, primary_ctx=primary_ctx)
+                    ctx = CudaContext(
+                        device_id=device_id,
+                        primary_ctx=primary_ctx,
+                        current_ctx=current_ctx,
+                    )
                 else:
                     if not isinstance(cuda_context, CudaContext):
                         raise TypeError("cuda_context must be a CudaContext")
@@ -1624,6 +1647,10 @@ def from_dlpack(
                         raise ValueError(
                             "cuda_context.primary_ctx does not match primary_ctx"
                         )
+                    if bool(cuda_context.current_ctx) != bool(current_ctx):
+                        raise ValueError(
+                            "cuda_context.current_ctx does not match current_ctx"
+                        )
                     ctx = cython.cast(CudaContext, cuda_context)
 
                 frames_ref = ctx.get_frames_ctx(sw_fmt, width, height)

av/video/frame.pyi

@@ -34,7 +34,14 @@ class CudaContext:
     def device_id(self) -> int: ...
     @property
     def primary_ctx(self) -> bool: ...
-    def __init__(self, device_id: int = 0, primary_ctx: bool = True) -> None: ...
+    @property
+    def current_ctx(self) -> bool: ...
+    def __init__(
+        self,
+        device_id: int = 0,
+        primary_ctx: bool = True,
+        current_ctx: bool = False,
+    ) -> None: ...
 
 class VideoFrame(Frame):
     format: VideoFormat
@@ -109,4 +116,5 @@ class VideoFrame(Frame):
         device_id: int | None = None,
         primary_ctx: bool = True,
         cuda_context: CudaContext | None = None,
+        current_ctx: bool = False,
     ) -> VideoFrame: ...

docs/api/video.rst

@@ -33,6 +33,9 @@ Video Frames
 
 .. automodule:: av.video.frame
 
+.. autoclass:: CudaContext
+    :members:
+
 .. autoclass:: VideoFrame
 
     A single video frame.
@@ -80,6 +83,7 @@ Conversions
 
 .. automethod:: VideoFrame.from_image
 .. automethod:: VideoFrame.from_ndarray
+.. automethod:: VideoFrame.from_dlpack
 
 
 

tests/test_dlpack.py

@@ -531,6 +531,19 @@ def test_video_frame_from_dlpack_invalid_plane_object_raises_typeerror() -> None
         VideoFrame.from_dlpack((object(), object()), format="nv12", width=64, height=48)
 
 
+def test_cuda_context_flags() -> None:
+    default_ctx = av.video.frame.CudaContext()
+    assert default_ctx.primary_ctx is True
+    assert default_ctx.current_ctx is False
+
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        av.video.frame.CudaContext(primary_ctx=True, current_ctx=True)
+
+    ctx = av.video.frame.CudaContext(primary_ctx=False, current_ctx=True)
+    assert ctx.primary_ctx is False
+    assert ctx.current_ctx is True
+
+
 def test_video_frame_from_dlpack_cuda_hw_frame_behavior_if_available() -> None:
     backend = _get_cuda_backend()
     if backend is None:
@@ -621,7 +634,10 @@ def test_video_frame_from_dlpack_cuda_hw_frame_behavior_if_available() -> None:
         pytest.skip(f"CUDA hwcontext not available in this build/runtime: {e}")
 
 
-def test_encode_cuda_frame_with_nvenc_if_available() -> None:
+@pytest.mark.parametrize(
+    "use_current_ctx", [False, True], ids=["primary-context", "current-context"]
+)
+def test_encode_cuda_frame_with_nvenc_if_available(use_current_ctx: bool) -> None:
     # Issue #2199: a CUDA frame from DLPack should encode on the GPU directly.
     # Its hw_frames_ctx must propagate to the encoder before avcodec_open2.
     backend = _get_cuda_backend()
@@ -639,7 +655,21 @@ def test_encode_cuda_frame_with_nvenc_if_available() -> None:
             y = mod.zeros((height, width), dtype=mod.uint8)
             uv = mod.zeros((height // 2, width // 2, 2), dtype=mod.uint8)
 
-        frame = VideoFrame.from_dlpack((y, uv), format="nv12")
+        if use_current_ctx:
+            current_ctx = av.video.frame.CudaContext(
+                device_id=int(y.__dlpack_device__()[1]),
+                primary_ctx=False,
+                current_ctx=True,
+            )
+            frame = VideoFrame.from_dlpack(
+                (y, uv),
+                format="nv12",
+                primary_ctx=False,
+                cuda_context=current_ctx,
+                current_ctx=True,
+            )
+        else:
+            frame = VideoFrame.from_dlpack((y, uv), format="nv12")
         assert frame.format.name == "cuda"
         assert frame.sw_format is not None and frame.sw_format.name == "nv12"