Improve from_dlpack CUDA context ergonomics

CHANGELOG.rst

@@ -28,6 +28,14 @@ We are operating with `semantic versioning <https://semver.org>`_.
     - Bug fixes (PATCH) go here. 
     - $CHANGE by :gh-user:`mikeboers` in (:pr:`1`).
 
+v18.1.0 (Unreleased)
+--------------------
+
+Features:
+
+- Support reusing the thread's current CUDA context via a ``current_ctx`` flag on ``CudaContext`` and ``VideoFrame.from_dlpack``, for interop with libraries like PyTorch that initialize CUDA first by :gh-user:`Yozer` (:pr:`2339`).
+- ``VideoFrame.from_dlpack`` no longer requires restating ``primary_ctx``/``current_ctx`` when passing an explicit ``cuda_context``; the flags are only validated when explicitly given by :gh-user:`WyattBlue`.
+
 v18.0.0
 -------
 

av/video/frame.py

@@ -1487,9 +1487,9 @@ def from_dlpack(
         height: int = 0,
         stream=None,
         device_id: int | None = None,
-        primary_ctx: bool = True,
+        primary_ctx: bool | None = None,
         cuda_context=None,
-        current_ctx: bool = False,
+        current_ctx: bool | None = None,
     ):
         if not isinstance(planes, (tuple, list)):
             planes = (planes,)
@@ -1633,8 +1633,10 @@ def from_dlpack(
                 if cuda_context is None:
                     ctx = CudaContext(
                         device_id=device_id,
-                        primary_ctx=primary_ctx,
-                        current_ctx=current_ctx,
+                        primary_ctx=(
+                            not current_ctx if primary_ctx is None else primary_ctx
+                        ),
+                        current_ctx=bool(current_ctx),
                     )
                 else:
                     if not isinstance(cuda_context, CudaContext):
@@ -1643,11 +1645,15 @@ def from_dlpack(
                         raise ValueError(
                             "cuda_context.device_id does not match the DLPack tensor device_id"
                         )
-                    if bool(cuda_context.primary_ctx) != bool(primary_ctx):
+                    if primary_ctx is not None and bool(
+                        cuda_context.primary_ctx
+                    ) != bool(primary_ctx):
                         raise ValueError(
                             "cuda_context.primary_ctx does not match primary_ctx"
                         )
-                    if bool(cuda_context.current_ctx) != bool(current_ctx):
+                    if current_ctx is not None and bool(
+                        cuda_context.current_ctx
+                    ) != bool(current_ctx):
                         raise ValueError(
                             "cuda_context.current_ctx does not match current_ctx"
                         )

av/video/frame.pyi

@@ -114,7 +114,7 @@ class VideoFrame(Frame):
         height: int = 0,
         stream: int | None = None,
         device_id: int | None = None,
-        primary_ctx: bool = True,
+        primary_ctx: bool | None = None,
         cuda_context: CudaContext | None = None,
-        current_ctx: bool = False,
+        current_ctx: bool | None = None,
     ) -> VideoFrame: ...

tests/test_dlpack.py

@@ -662,11 +662,7 @@ def test_encode_cuda_frame_with_nvenc_if_available(use_current_ctx: bool) -> Non
                 current_ctx=True,
             )
             frame = VideoFrame.from_dlpack(
-                (y, uv),
-                format="nv12",
-                primary_ctx=False,
-                cuda_context=current_ctx,
-                current_ctx=True,
+                (y, uv), format="nv12", cuda_context=current_ctx
             )
         else:
             frame = VideoFrame.from_dlpack((y, uv), format="nv12")