diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 0c86d12e4..8dabf3fcb 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -28,6 +28,14 @@ We are operating with `semantic versioning `_. - 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 ------- diff --git a/av/video/frame.py b/av/video/frame.py index 511cb028b..4a85d5b11 100644 --- a/av/video/frame.py +++ b/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" ) diff --git a/av/video/frame.pyi b/av/video/frame.pyi index 14197f24f..dd6f481d1 100644 --- a/av/video/frame.pyi +++ b/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: ... diff --git a/tests/test_dlpack.py b/tests/test_dlpack.py index cfff0540c..033316138 100644 --- a/tests/test_dlpack.py +++ b/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")