feat(nvenc): add support for explicit cuda streams

av/video/frame.pxd

@@ -1,5 +1,5 @@
 cimport libav as lib
-from libc.stdint cimport uint8_t
+from libc.stdint cimport uint8_t, uintptr_t
 
 from av.frame cimport Frame
 from av.video.format cimport VideoFormat
@@ -10,6 +10,7 @@ cdef class CudaContext:
     cdef readonly int device_id
     cdef readonly bint primary_ctx
     cdef readonly bint current_ctx
+    cdef readonly uintptr_t cuda_stream
     cdef lib.AVBufferRef* _device_ref
     cdef dict _frames_cache
     cdef lib.AVBufferRef* _get_device_ref(self)

av/video/frame.py

@@ -16,7 +16,8 @@
     PyCapsule_SetName,
 )
 from cython.cimports.cpython.ref import Py_DECREF, Py_INCREF
-from cython.cimports.libc.stdint import int64_t, uint8_t
+from cython.cimports.hwcontext_cuda import AVCUDADeviceContext, CUstream
+from cython.cimports.libc.stdint import int64_t, uint8_t, uintptr_t
 
 
 @cython.final
@@ -29,21 +30,32 @@ class CudaContext:
     :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``.
+    :param int cuda_stream: Optional raw ``CUstream`` pointer for FFmpeg CUDA
+        operations, including NVENC input and output. The caller owns the
+        stream and must keep it alive as long as this context can be used.
     """
 
     def __cinit__(
         self,
         device_id: cython.int = 0,
         primary_ctx: cython.bint = True,
         current_ctx: cython.bint = False,
+        cuda_stream: object = None,
     ):
         self.device_id = device_id
         self.primary_ctx = primary_ctx
         self.current_ctx = current_ctx
+        self.cuda_stream = 0
         self._device_ref = cython.NULL
         self._frames_cache = {}
         if primary_ctx and current_ctx:
             raise ValueError("primary_ctx and current_ctx are mutually exclusive")
+        if cuda_stream is not None:
+            if not isinstance(cuda_stream, int):
+                raise TypeError("cuda_stream must be an integer or None")
+            if cuda_stream < 0:
+                raise ValueError("cuda_stream must be non-negative")
+            self.cuda_stream = cython.cast(uintptr_t, cuda_stream)
 
     def __dealloc__(self):
         ref: cython.pointer[lib.AVBufferRef]
@@ -83,6 +95,14 @@ def _get_device_ref(self) -> cython.pointer[lib.AVBufferRef]:
                 0,
             )
         )
+        if self.cuda_stream:
+            device_ctx: cython.pointer[lib.AVHWDeviceContext] = cython.cast(
+                cython.pointer[lib.AVHWDeviceContext], device_ref.data
+            )
+            cuda_ctx: cython.pointer[AVCUDADeviceContext] = cython.cast(
+                cython.pointer[AVCUDADeviceContext], device_ctx.hwctx
+            )
+            cuda_ctx.stream = cython.cast(CUstream, self.cuda_stream)
         self._device_ref = device_ref
         return device_ref
 

av/video/frame.pyi

@@ -36,11 +36,14 @@ class CudaContext:
     def primary_ctx(self) -> bool: ...
     @property
     def current_ctx(self) -> bool: ...
+    @property
+    def cuda_stream(self) -> int: ...
     def __init__(
         self,
         device_id: int = 0,
         primary_ctx: bool = True,
         current_ctx: bool = False,
+        cuda_stream: int | None = None,
     ) -> None: ...
 
 class VideoFrame(Frame):

include/avutil.pxd

@@ -192,7 +192,7 @@ cdef extern from "libavutil/frame.h" nogil:
 
 cdef extern from "libavutil/hwcontext.h" nogil:
     cdef struct AVHWDeviceContext:
-        pass
+        void *hwctx
 
     enum AVHWDeviceType:
         AV_HWDEVICE_TYPE_NONE

include/hwcontext_cuda.pxd

@@ -0,0 +1,17 @@
+cdef extern from *:
+    """
+    #ifndef CUDA_VERSION
+    #define PYAV_CUDA_TYPES
+    #define CUDA_VERSION 1
+    typedef struct CUctx_st *CUcontext;
+    typedef struct CUstream_st *CUstream;
+    #endif
+    #include <libavutil/hwcontext_cuda.h>
+    #ifdef PYAV_CUDA_TYPES
+    #undef CUDA_VERSION
+    #undef PYAV_CUDA_TYPES
+    #endif
+    """
+    ctypedef void *CUstream
+    ctypedef struct AVCUDADeviceContext:
+        CUstream stream

tests/test_dlpack.py

@@ -533,13 +533,23 @@ 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
+    assert default_ctx.cuda_stream == 0
 
     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)
+    ctx = av.video.frame.CudaContext(
+        primary_ctx=False, current_ctx=True, cuda_stream=1234
+    )
     assert ctx.primary_ctx is False
     assert ctx.current_ctx is True
+    assert ctx.cuda_stream == 1234
+
+    with pytest.raises(TypeError, match="integer or None"):
+        av.video.frame.CudaContext(cuda_stream="1234")  # type: ignore[arg-type]
+
+    with pytest.raises(ValueError, match="non-negative"):
+        av.video.frame.CudaContext(cuda_stream=-1)
 
 
 def test_video_frame_from_dlpack_cuda_hw_frame_behavior_if_available() -> None:
@@ -679,3 +689,52 @@ def test_encode_cuda_frame_with_nvenc_if_available(use_current_ctx: bool) -> Non
         assert any(p.size for p in packets)
     except av.FFmpegError as e:
         pytest.skip(f"nvenc/CUDA not available in this build/runtime: {e}")
+
+
+def test_encode_cuda_frame_with_nvenc_external_stream_if_available() -> None:
+    try:
+        import torch  # type: ignore
+    except Exception:
+        pytest.skip("PyTorch is not available.")
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is not available.")
+
+    width, height = 256, 256
+    producer_stream = torch.cuda.Stream()
+    encoder_stream = torch.cuda.Stream()
+    encoder_stream_ptr = int(encoder_stream.cuda_stream)
+
+    try:
+        with torch.cuda.stream(producer_stream):
+            y = torch.zeros((height, width), dtype=torch.uint8, device="cuda")
+            uv = torch.zeros(
+                (height // 2, width // 2, 2),
+                dtype=torch.uint8,
+                device="cuda",
+            )
+            cuda_context = av.video.frame.CudaContext(
+                device_id=int(y.__dlpack_device__()[1]),
+                primary_ctx=False,
+                current_ctx=True,
+                cuda_stream=encoder_stream_ptr,
+            )
+            frame = VideoFrame.from_dlpack(
+                (y, uv),
+                format="nv12",
+                stream=encoder_stream_ptr,
+                cuda_context=cuda_context,
+            )
+
+        assert cuda_context.cuda_stream == encoder_stream_ptr
+        cc = av.CodecContext.create("h264_nvenc", "w")
+        cc.width = width
+        cc.height = height
+        cc.time_base = Fraction(1, 24)
+        cc.framerate = Fraction(24, 1)
+        cc.pix_fmt = "cuda"
+
+        packets = cc.encode(frame)
+        packets += cc.encode(None)
+        assert any(p.size for p in packets)
+    except av.FFmpegError as e:
+        pytest.skip(f"nvenc/CUDA not available in this build/runtime: {e}")