Preserving hardware memory during cuvid decoding, exporting/importing via dlpack.

[caffeinism]

#2148

Hello? I'm a user with limited knowledge of libav, dlpack, and cython. However, recognizing this as a necessary feature, I drafted this with the help of an LLM.

Motivation

If an application decodes video, performs GPU operations, and then re-encodes it, PyAV currently incurs a significant amount of memcopy. (GPU (cuvid) -> CPU (PyAV) -> GPU (Torch, etc.) -> CPU (PyAV) -> GPU (nvenc)) However, if we could export frames decoded by cuvid to dlpack while keeping them on the GPU, we wouldn't need to move the frames to CPU memory.

I passed all existing tests, but with such extensive modifications, it seems difficult for a beginner like me to catch every single detail. However, since most changes involve adding features rather than modifying existing ones, I hope this PR serves as a good starting point.

Usage example

import av
from av.codec.hwaccel import HWAccel
import torch

hwaccel = HWAccel(
    device_type="cuda",
    device=0,
    allow_software_fallback=False,
    output_format="hw", # preserve hw memory
)

# decode using cuvid
with av.open(from_video_filename, "r", hwaccel=hwaccel) as c:
    frame = next(c.decode(video=0))
    y = torch.from_dlpack(frame.planes[0]) # device(type='cuda', index=0), torch.uint8, torch.Size([H, W])
    uv = torch.from_dlpack(frame.planes[1]) # device(type='cuda', index=0), torch.uint8, torch.Size([H/2, W/2])

f = av.VideoFrame.from_dlpack(((y*0.5).to(torch.uint8), uv)) # some operation

with av.open(to_video_filename, "w") as c:
    s = c.add_stream("h264_nvenc", rate=24) # encode using nvenc
    for it in s.encode(f):
        c.mux(it)
    for it in s.encode(None):
        c.mux(it)

[caffeinism]

@WyattBlue If I add tests, will it work fine even if it only runs on a CUDA machine? I don't think it will work in the GitHub workflow.

[WyattBlue]

You need to test the interface. For example, hw_format does not have an pyi interface, and writing a test would catch that fact.

[WyattBlue]

av/hwcontext.pxd‎ should be merged with include/avutil. *.pxd files should otherwise not be free radicals, i.e., they should have a corresponding real .py file.

[caffeinism]

You need to test the interface. For example, hw_format does not have an pyi interface, and writing a test would catch that fact.

Could you please explain it in a bit more detail?

av/hwcontext.pxd‎ should be merged with include/avutil. *.pxd files should otherwise not be free radicals, i.e., they should have a corresponding real .py file.

In this case, how should dlpack.pxd be handled? Should this also be moved to the include directory?

[caffeinism]

@WyattBlue Could you take a look at the last commit section? I modified the buffer creation logic for frames generated by VideoFrame(), VideoFrame.from_ndarray(), and VideoFrame.reformat to support dlpack.