Add support for NVDEC backend in video processing and enhance README documentation. Introduce a new script for extracting EasyCom sessions on specific GPUs, detailing session management and error handling. Update requirements.txt to include PyNvVideoCodec as an optional dependency. Refactor streaming_turn_model_directly.py to accommodate multiple decoder backends and improve frame indexing with ffprobe.

streaming_turn_model_directly.py

@@ -95,7 +95,7 @@ class IndexedSourceClip:
     dense_end: float
     duration: float
     crop: tuple[int, int, int, int] | None
-    reader: Any
+    reader: Any | None
     source_frame_starts: np.ndarray
     source_frame_ends: np.ndarray
     source_fps: float
@@ -151,6 +151,76 @@ def _timestamp_array(reader: Any) -> tuple[np.ndarray, np.ndarray]:
     return starts, ends
 
 
+def _parse_ffprobe_rate(value: str, path: Path) -> float:
+    """Parse an ffprobe rational such as ``20/1``."""
+
+    try:
+        numerator, denominator = value.split("/", maxsplit=1)
+        rate = float(numerator) / float(denominator)
+    except (AttributeError, ValueError, ZeroDivisionError) as exc:
+        raise RuntimeError(f"Invalid ffprobe frame rate {value!r}: {path}") from exc
+    if not math.isfinite(rate) or rate <= 0:
+        raise RuntimeError(f"Non-positive ffprobe frame rate {rate}: {path}")
+    return rate
+
+
+def _ffprobe_video_index(
+    path: Path,
+) -> tuple[np.ndarray, np.ndarray, float, int, int]:
+    """Index display-order frame PTS without keeping a software decoder open."""
+
+    if shutil.which("ffprobe") is None:
+        raise RuntimeError("ffprobe is required by the pynvcodec backend")
+    completed = subprocess.run(
+        [
+            "ffprobe",
+            "-v",
+            "error",
+            "-select_streams",
+            "v:0",
+            "-show_frames",
+            "-show_entries",
+            (
+                "stream=width,height,avg_frame_rate:"
+                "frame=best_effort_timestamp_time,pkt_duration_time"
+            ),
+            "-of",
+            "json",
+            str(path),
+        ],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+    if completed.returncode != 0:
+        raise RuntimeError(
+            f"Failed to index video with ffprobe: {path}: "
+            f"{completed.stderr.strip()}"
+        )
+    try:
+        payload = json.loads(completed.stdout)
+        stream = payload["streams"][0]
+        width = int(stream["width"])
+        height = int(stream["height"])
+        source_fps = _parse_ffprobe_rate(stream["avg_frame_rate"], path)
+        raw_starts = [
+            float(frame["best_effort_timestamp_time"])
+            for frame in payload["frames"]
+        ]
+    except (IndexError, KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
+        raise RuntimeError(f"Incomplete ffprobe frame index for {path}") from exc
+
+    if not raw_starts:
+        raise RuntimeError(f"Video contains no indexed frames: {path}")
+    starts = np.asarray(raw_starts, dtype=np.float64)
+    starts -= float(starts[0])
+    if np.any(np.diff(starts) < -1e-9):
+        raise RuntimeError(f"ffprobe video frame PTS values are not monotonic: {path}")
+    nominal = 1.0 / source_fps
+    ends = np.concatenate((starts[1:], starts[-1:] + nominal))
+    return starts, ends, source_fps, width, height
+
+
 def _decode_audio_exact(path: Path, duration: float) -> np.ndarray:
     """Decode, then pad/trim one clip to its dense video duration."""
 
@@ -203,6 +273,8 @@ class CausalSourceAV:
         output_fps: float,
         required_duration: float,
         frame_cache_size: int = 8,
+        decoder_backend: str = "decord",
+        decoder_gpu_id: int = 0,
     ) -> None:
         if not specs:
             raise ValueError("At least one source clip is required")
@@ -212,20 +284,40 @@ class CausalSourceAV:
             raise ValueError("required_duration must be positive")
         if frame_cache_size < 0:
             raise ValueError("frame_cache_size must be non-negative")
+        if decoder_backend not in {"decord", "pynvcodec"}:
+            raise ValueError(f"Unsupported decoder backend: {decoder_backend}")
+        if decoder_gpu_id < 0:
+            raise ValueError("decoder_gpu_id must be non-negative")
         if shutil.which("ffmpeg") is None:
             raise RuntimeError("ffmpeg is required to decode source audio")
 
-        try:
-            import decord
-        except ImportError as exc:
-            raise RuntimeError(
-                "Direct extraction requires decord; install "
-                "qwen-omni-utils[decord]"
-            ) from exc
+        decord = None
+        self._nvc = None
+        if decoder_backend == "decord":
+            try:
+                import decord
+            except ImportError as exc:
+                raise RuntimeError(
+                    "The decord backend requires qwen-omni-utils[decord]"
+                ) from exc
+        else:
+            try:
+                import PyNvVideoCodec as nvc
+            except ImportError as exc:
+                raise RuntimeError(
+                    "The pynvcodec backend requires PyNvVideoCodec; install it "
+                    "with `python -m pip install PyNvVideoCodec`"
+                ) from exc
+            self._nvc = nvc
 
         self.fps = float(output_fps)
         self.required_duration = float(required_duration)
         self.frame_cache_size = int(frame_cache_size)
+        self.decoder_backend = decoder_backend
+        self.decoder_gpu_id = int(decoder_gpu_id)
+        self._nvc_decoder = None
+        self._nvc_source_path: Path | None = None
+        self._nvc_decoder_reconfigurations = 0
         self._frame_cache: OrderedDict[tuple[int, int], Any] = OrderedDict()
         self._request_count = 0
         self._cache_hits = 0
@@ -246,19 +338,24 @@ class CausalSourceAV:
                 f"{path.name}",
                 flush=True,
             )
-            reader = decord.VideoReader(str(path))
-            if not len(reader):
-                raise RuntimeError(f"No video frames decoded from {path}")
-            starts, ends = _timestamp_array(reader)
-            source_fps = float(reader.get_avg_fps())
-            first_shape = tuple(int(value) for value in reader[0].shape)
-            if len(first_shape) != 3:
-                raise RuntimeError(
-                    f"Unexpected decoded frame shape {first_shape}: {path}"
-                )
-            height, width = first_shape[:2]
+            if decoder_backend == "decord":
+                assert decord is not None
+                reader = decord.VideoReader(str(path))
+                if not len(reader):
+                    raise RuntimeError(f"No video frames decoded from {path}")
+                starts, ends = _timestamp_array(reader)
+                source_fps = float(reader.get_avg_fps())
+                first_shape = tuple(int(value) for value in reader[0].shape)
+                if len(first_shape) != 3:
+                    raise RuntimeError(
+                        f"Unexpected decoded frame shape {first_shape}: {path}"
+                    )
+                height, width = first_shape[:2]
+            else:
+                reader = None
+                starts, ends, source_fps, width, height = _ffprobe_video_index(path)
 
-            inferred_duration = len(reader) / source_fps
+            inferred_duration = len(starts) / source_fps
             duration = (
                 float(spec.duration)
                 if spec.duration is not None
@@ -389,7 +486,25 @@ class CausalSourceAV:
             return cached.copy()
 
         clip = self.clips[clip_index]
-        array = clip.reader[frame_index].asnumpy()
+        if self.decoder_backend == "decord":
+            if clip.reader is None:
+                raise RuntimeError(f"Missing decord reader for {clip.path}")
+            array = clip.reader[frame_index].asnumpy()
+        else:
+            decoder = self._pynv_decoder_for_clip(clip)
+            decoded_frame = decoder[frame_index]
+            frame_tensor = torch.from_dlpack(decoded_frame)
+            if frame_tensor.ndim != 3 or frame_tensor.shape[-1] != 3:
+                raise RuntimeError(
+                    "PyNvVideoCodec RGB output must have HWC shape, got "
+                    f"{tuple(frame_tensor.shape)} for {clip.path}"
+                )
+            if frame_tensor.dtype != torch.uint8:
+                raise RuntimeError(
+                    "PyNvVideoCodec RGB output must be uint8, got "
+                    f"{frame_tensor.dtype} for {clip.path}"
+                )
+            array = frame_tensor.cpu().contiguous().numpy()
         if clip.crop is not None:
             x, y, width, height = clip.crop
             array = array[y : y + height, x : x + width]
@@ -401,6 +516,44 @@ class CausalSourceAV:
                 self._frame_cache.popitem(last=False)
         return image.copy()
 
+    def _pynv_decoder_for_clip(self, clip: IndexedSourceClip) -> Any:
+        """Return one reusable NVDEC decoder configured for ``clip``."""
+
+        if self._nvc is None:
+            raise RuntimeError("PyNvVideoCodec was not initialized")
+        source_path = clip.path
+        if self._nvc_decoder is None:
+            self._nvc_decoder = self._new_pynv_decoder(source_path)
+            self._nvc_source_path = source_path
+        elif self._nvc_source_path != source_path:
+            torch.cuda.synchronize(self.decoder_gpu_id)
+            try:
+                self._nvc_decoder.reconfigure_decoder(str(source_path))
+            except Exception:
+                # A rebuild is safer than silently using stale decoder state.
+                self._nvc_decoder = self._new_pynv_decoder(source_path)
+            self._nvc_source_path = source_path
+            self._nvc_decoder_reconfigurations += 1
+
+        decoded_length = len(self._nvc_decoder)
+        if decoded_length != clip.num_frames:
+            raise RuntimeError(
+                "PyNvVideoCodec/ffprobe frame-count mismatch for "
+                f"{clip.path}: decoder={decoded_length}, ffprobe={clip.num_frames}"
+            )
+        return self._nvc_decoder
+
+    def _new_pynv_decoder(self, source_path: Path) -> Any:
+        assert self._nvc is not None
+        return self._nvc.SimpleDecoder(
+            str(source_path),
+            gpu_id=self.decoder_gpu_id,
+            use_device_memory=True,
+            need_scanned_stream_metadata=True,
+            decoder_cache_size=4,
+            output_color_type=self._nvc.OutputColorType.RGB,
+        )
+
     def frames(self, start: float, end: float) -> list[Any]:
         """Return causal global-grid frames in ``[start, end)``."""
 
@@ -451,7 +604,20 @@ class CausalSourceAV:
 
     def metadata(self) -> dict[str, Any]:
         return {
-            "decoder": "decord_with_frame_pts_and_ffmpeg_pcm_audio",
+            "decoder": (
+                "decord_with_frame_pts_and_ffmpeg_pcm_audio"
+                if self.decoder_backend == "decord"
+                else "pynvvideocodec_nvdec_with_ffprobe_pts_and_ffmpeg_pcm_audio"
+            ),
+            "decoder_backend": self.decoder_backend,
+            "decoder_gpu_id": (
+                self.decoder_gpu_id if self.decoder_backend == "pynvcodec" else None
+            ),
+            "decoder_version": (
+                getattr(self._nvc, "__version__", None)
+                if self.decoder_backend == "pynvcodec"
+                else None
+            ),
             "visual_sampling": "global_grid_latest_source_pts_not_after_tick",
             "interval_convention": "half_open_[start,end)",
             "output_video_fps": self.fps,
@@ -483,6 +649,7 @@ class CausalSourceAV:
                 "odd_frame_duplications": self._odd_frame_duplications,
                 "short_interval_fallbacks": self._short_interval_fallbacks,
                 "tail_frame_reuses": self._tail_frame_reuses,
+                "decoder_reconfigurations": self._nvc_decoder_reconfigurations,
             },
         }
 
@@ -609,6 +776,18 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28)
     parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28)
     parser.add_argument("--frame-cache-size", type=int, default=8)
+    parser.add_argument(
+        "--decoder",
+        choices=("decord", "pynvcodec"),
+        default="decord",
+        help="Video pixel decoder; pynvcodec uses NVIDIA NVDEC",
+    )
+    parser.add_argument(
+        "--decoder-gpu-id",
+        type=int,
+        default=0,
+        help="GPU visible to this process used by PyNvVideoCodec",
+    )
     parser.add_argument("--verify-every", type=int, default=0)
     parser.add_argument("--verify-continuous-every", type=int, default=0)
     parser.add_argument("--no-progress", action="store_true")
@@ -742,6 +921,8 @@ def main() -> None:
         output_fps=args.video_fps,
         required_duration=required_duration,
         frame_cache_size=args.frame_cache_size,
+        decoder_backend=args.decoder,
+        decoder_gpu_id=args.decoder_gpu_id,
     )
     is_easycom = (
         args.easycom_session_dir is not None
@@ -772,6 +953,10 @@ def main() -> None:
                 )
     _print_media_summary(reader, args.start, end)
     if args.media_only:
+        # Exercise the selected pixel decoder as well as its metadata path.
+        probe_end = min(end, args.start + max(args.step_seconds, 1.0 / reader.fps))
+        reader.frames(args.start, probe_end)
+        print(f"Decoded media probe successfully with {args.decoder}")
         return
 
     duration = end - args.start