nvidia-codec 0.0.20__tar.gz → 0.0.21__tar.gz

{nvidia_codec-0.0.20/src/nvidia_codec.egg-info → nvidia_codec-0.0.21}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nvidia-codec
-Version: 0.0.20
+Version: 0.0.21
 Summary: Pythonic interface to Nvidia Codec
 Home-page: https://github.com/riaqn/python-nvidia-codec
 Author: Zesen Qian

{nvidia_codec-0.0.20 → nvidia_codec-0.0.21}/setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = nvidia-codec
-version = 0.0.20
+version = 0.0.21
 author = Zesen Qian
 author_email = nvidia-codec@riaqn.org
 description = Pythonic interface to Nvidia Codec

{nvidia_codec-0.0.20 → nvidia_codec-0.0.21}/src/nvidia_codec/core/decode.py

@@ -17,6 +17,7 @@ from .common import *
 from .. import CodecNotSupportedError
 import numpy as np
 from ctypes import *
+import weakref
 
 import logging
 
@@ -123,12 +124,19 @@ class Surface:
 class Picture:
     """Decoded frame before post-processing.
 
-    A data object — does not own the picture slot. Slot ownership is
-    managed by the caller via pre_decode/pic_release callbacks.
+    A pure data object — index, proc params, pts. Deliberately holds NO
+    reference to the decoder: the decoder tracks live Pictures in
+    _slots/_recent, and a back-reference here would form a cycle that
+    keeps the decoder (and its large NVDEC context) alive past the
+    refcount-zero point, deferring cleanup to the cyclic GC. To turn a
+    Picture into a Surface use Decoder.map(pic) (the decoder is always
+    in scope where mapping happens).
+
+    Slot ownership is managed by the caller via pre_decode/pic_release
+    callbacks.
     """
 
-    def __init__(self, decoder, index, proc_params, pts):
-        self.decoder = decoder
+    def __init__(self, index, proc_params, pts):
         self.index = index
         self.params = proc_params
         self.pts = pts
@@ -137,9 +145,6 @@ class Picture:
         """No-op. Slot ownership is managed by the caller (player.py)."""
         pass
 
-    def map(self, stream: int = 0):
-        return Surface(self.decoder, self.index, self.params, stream, self.pts)
-
 
 def decide_surface_format(
     chroma_format, bit_depth, supported_surface_formats, allow_high=False
@@ -192,12 +197,27 @@ class BaseDecoder:
     """
 
     def catch_exception(self, func, return_on_error=0):
+        # `func` is a bound method of self, and the wrapper records errors
+        # on self.exception — both strong refs to self. The wrapper is
+        # wrapped in a CFUNCTYPE and stored on self (self.handleXxxCallback),
+        # so capturing self strongly here forms a self-cycle that makes the
+        # decoder uncollectable by refcounting (only the cyclic GC can take
+        # it — which under streaming load lags and balloons VRAM). Hold
+        # self weakly instead; the callback only ever fires during a
+        # send()/cuvidParseVideoData while a caller keeps the decoder
+        # alive, so the weakref always resolves then. The None branch is
+        # just belt-and-suspenders for a callback racing finalization.
+        selfref = weakref.ref(self)
+        fn = func.__func__
         def wrapper(*args, **kwargs):
+            s = selfref()
+            if s is None:
+                return return_on_error
             try:
-                return func(*args, **kwargs)
+                return fn(s, *args, **kwargs)
             except BaseException as e:
                 log.debug(f"callback exception logged {e}")
-                self.exception = e
+                s.exception = e
                 log.debug(f"returning error code {return_on_error}")
                 return return_on_error
 
@@ -410,10 +430,16 @@ class BaseDecoder:
             unpaired_field=di.repeat_first_field < 0,
         )
 
-        picture = Picture(self, di.picture_index, params, di.timestamp)
+        picture = Picture(di.picture_index, params, di.timestamp)
         self.post_decode(picture)
         return 1
 
+    def map(self, pic, stream: int = 0):
+        """Map a decoded Picture to a Surface (post-processed frame).
+        Lives on the decoder (not Picture.map) so Picture needn't hold a
+        decoder back-reference — see Picture's docstring."""
+        return Surface(self, pic.index, pic.params, stream, pic.pts)
+
     def handleOperatingPoint(self, pUserData, pOPInfo):
         opi = pOPInfo.contents
         if opi.codec == cudaVideoCodec.AV1:
@@ -528,18 +554,41 @@ class BaseDecoder:
                 payload=data.ctypes.data_as(POINTER(c_uint8)),
                 timestamp=pts,
             )
+        # self.exception relays an error out of a ctypes callback (which
+        # can't propagate through the C parser). Clear it as we consume
+        # it: a retained exception keeps its __traceback__ alive, whose
+        # frames reference self, pinning the decoder. Raise it WITHOUT
+        # binding it to a local here — a `e = self.exception; raise e`
+        # would put `e` in this frame, and the raised exception's
+        # traceback includes this same frame → e ↔ frame cycle, defeating
+        # the point. `try: raise self.exception finally: clear` keeps the
+        # exception out of this frame's locals.
         self.exception = None
-        with cuda.Device(self.device):
-            try:
+        try:
+            with cuda.Device(self.device):
                 cuda.check(nvcuvid.cuvidParseVideoData(self.cuvid_parser, byref(p)))
-            except cuda.CUError:
-                if self.exception is not None:
-                    raise self.exception
+        except cuda.CUError:
+            if self.exception is None:
                 raise
+            # else: a callback captured a more specific exception — it's
+            # raised (and cleared) below in preference to the CUError.
         if self.exception is not None:
-            raise self.exception
+            try:
+                raise self.exception
+            finally:
+                self.exception = None
 
     def free(self):
+        # __del__ can call this during interpreter shutdown, by which
+        # point module globals (nvcuvid, cuda, the ctypes types) may be
+        # torn down to None — calling through them then raises
+        # "NoneType is not callable". The OS reclaims the GPU context on
+        # process exit regardless, so just bail if our entry points are
+        # gone. (Now reachable because the decoder is refcount-collectable
+        # — its self-reference cycles were removed — so __del__ actually
+        # runs at shutdown instead of the object leaking.)
+        if nvcuvid is None or cuda is None:
+            return
         if getattr(self, 'cuvid_parser', None):
             cuda.check(nvcuvid.cuvidDestroyVideoParser(self.cuvid_parser))
             self.cuvid_parser = CUvideoparser()
@@ -549,7 +598,10 @@ class BaseDecoder:
                 self.cuvid_decoder = CUvideodecoder()
 
     def __del__(self):
-        self.free()
+        try:
+            self.free()
+        except Exception:
+            pass   # best-effort during finalization / shutdown
 
     def __enter__(self):
         return self

{nvidia_codec-0.0.20 → nvidia_codec-0.0.21}/src/nvidia_codec/ffmpeg/include/libavcodec.py

@@ -15,6 +15,10 @@ class AVMediaType(TypedCEnumeration(c_int)):
     NB = 5
 
 class AVCodecID(TypedCEnumeration(c_int)):
+    # Values track the ffmpeg 9 enum (libavcodec 63, pinned via the versioned
+    # libavcodec.so.63 load in ffmpeg/libavcodec.py — ffmpeg 8 removed entries,
+    # shifting WEBP/VP9/HEVC/HQ_HQA/AV1 down). A mismatch surfaces as
+    # CodecNotSupportedError ("<AVCodecID.(unknown): N>") on NVDEC-supported files.
     MPEG1 = 1
     MPEG2 = 2
     MPEG4 = 12
@@ -30,11 +34,12 @@ class AVCodecID(TypedCEnumeration(c_int)):
     VP6F = 92
     GIF = 97
     PRORES = 147
-    WEBP = 171
+    WEBP = 170
     VP8 = 139
-    VP9 = 167
-    HEVC = 173
-    AV1 = 225
+    VP9 = 166
+    HEVC = 172
+    HQ_HQA = 185
+    AV1 = 222
     # Audio codecs — values must match ffmpeg's avcodec.h
     PCM_S16LE = 0x10000
     PCM_S16BE = 0x10001

{nvidia_codec-0.0.20 → nvidia_codec-0.0.21}/src/nvidia_codec/ffmpeg/libavcodec.py

@@ -8,7 +8,10 @@ from errno import EAGAIN
 import  logging
 log = logging.getLogger(__name__)
 
-lib = cdll.LoadLibrary('libavcodec.so')
+# Pin the ABI major: enum values + struct layouts in include/libavcodec.py
+# track ffmpeg 9 (libavcodec 63). Loading the versioned .so makes a system
+# ffmpeg upgrade fail loudly here instead of mis-decoding with shifted enums.
+lib = cdll.LoadLibrary('libavcodec.so.63')
 
 
 av_packet_alloc = lib.av_packet_alloc

{nvidia_codec-0.0.20 → nvidia_codec-0.0.21}/src/nvidia_codec/utils/color.py

@@ -23,6 +23,44 @@ from ..ffmpeg.include.libavutil import AVColorRange, AVColorSpace
 from ..core.cuviddec import cudaVideoSurfaceFormat
 from ..core.decode import Surface
 
+def to_yuv420p(surface : Surface):
+    """De-interleave/subsample a decoded surface into a contiguous planar I420
+    (yuv420p) uint8 CUDA tensor: Y (H·W) ++ U (H·W/4) ++ V (H·W/4).
+
+    Unlike convert() — which applies the YUV→RGB color matrix — this does NO
+    color math, only plane layout: for NV12 (semi-planar 4:2:0) it splits the
+    interleaved chroma; for YUV444 (planar 4:4:4) it box-averages the full-res
+    chroma 2×2 down to 4:2:0. 8-bit only (NV12 / YUV444). Half the bytes of rgb24
+    and no matrix round-trip; the encoder consumes yuv420p directly. Returns a 1D
+    tensor of length H·W·3//2, ready for `.cpu().numpy().tobytes()`."""
+    fmt = surface.format
+    width, height = surface.width, surface.height
+    YUV = torch.as_tensor(surface, device=f'cuda:{surface.decoder.device}')
+    if fmt == cudaVideoSurfaceFormat.NV12:
+        Y = YUV[:height, :]          # (H, W)
+        UV = YUV[height:, :]         # (H/2, W): U,V interleaved
+        U = UV[:, 0::2]              # (H/2, W/2)
+        V = UV[:, 1::2]              # (H/2, W/2)
+    elif fmt == cudaVideoSurfaceFormat.YUV444:
+        # planar 4:4:4 → 4:2:0: keep Y full-res, box-average chroma 2×2. Plane
+        # layout mirrors convert() (3-plane tensor, or 3 planes stacked H·3).
+        if YUV.ndim == 3:
+            Y, U444, V444 = YUV[0], YUV[1], YUV[2]           # each (H, W)
+        else:
+            Y = YUV[:height, :]
+            U444 = YUV[height:height * 2, :]
+            V444 = YUV[height * 2:height * 3, :]
+        eh, ew = height // 2 * 2, width // 2 * 2
+        def _sub(c):                 # (H, W) 4:4:4 → (H/2, W/2) box-average, uint8
+            c = c[:eh, :ew].contiguous().to(torch.float32).reshape(eh // 2, 2, ew // 2, 2)
+            return c.mean(dim=(1, 3)).round().clamp(0, 255).to(torch.uint8)
+        U, V = _sub(U444), _sub(V444)
+    else:
+        raise ValueError(f"to_yuv420p: NV12/YUV444 8-bit only, got {fmt}")
+    # reshape(-1) copies the (possibly pitch-strided) planes to contiguous.
+    return torch.cat([Y.reshape(-1), U.reshape(-1), V.reshape(-1)])
+
+
 def convert(surface : Surface,
         source_space : AVColorSpace,
         source_range : AVColorRange,