Add TorchCodec video backend

cache_omni_av_pieces.py

@@ -95,7 +95,11 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
             "original maximum-length padding"
         ),
     )
-    parser.add_argument("--decoder", choices=("decord", "pynvcodec"), default="decord")
+    parser.add_argument(
+        "--decoder",
+        choices=("decord", "torchcodec", "pynvcodec"),
+        default="decord",
+    )
     parser.add_argument("--decoder-gpu-id", type=int, default=0)
     parser.add_argument("--decoder-threads", type=int, default=1)
     parser.add_argument("--max-open-readers", type=int, default=1)

extract_ego4d_batch.py

@@ -122,7 +122,7 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
     parser.add_argument("--frame-cache-size", type=int, default=64)
     parser.add_argument(
         "--decoder",
-        choices=("decord", "pynvcodec"),
+        choices=("decord", "torchcodec", "pynvcodec"),
         default="decord",
     )
     parser.add_argument("--decoder-gpu-id", type=int, default=0)

requirements.txt

@@ -12,4 +12,5 @@ scikit-learn
 peft
 
 # Optional on compatible NVIDIA GPUs: flash-attn
+# Optional CPU video backend: torchcodec
 # Optional NVDEC backend: PyNvVideoCodec

streaming_turn_model_directly.py

@@ -170,7 +170,7 @@ def _ffprobe_video_index(
     """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")
+        raise RuntimeError("ffprobe is required by non-Decord video backends")
     completed = subprocess.run(
         [
             "ffprobe",
@@ -288,7 +288,7 @@ 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"}:
+        if decoder_backend not in {"decord", "torchcodec", "pynvcodec"}:
             raise ValueError(f"Unsupported decoder backend: {decoder_backend}")
         if decoder_gpu_id < 0:
             raise ValueError("decoder_gpu_id must be non-negative")
@@ -301,6 +301,8 @@ class CausalSourceAV:
 
         decord = None
         self._decord = None
+        self._torchcodec = None
+        self._torchcodec_decoder_cls = None
         self._nvc = None
         if decoder_backend == "decord":
             try:
@@ -310,6 +312,17 @@ class CausalSourceAV:
                     "The decord backend requires qwen-omni-utils[decord]"
                 ) from exc
             self._decord = decord
+        elif decoder_backend == "torchcodec":
+            try:
+                import torchcodec
+                from torchcodec.decoders import VideoDecoder
+            except ImportError as exc:
+                raise RuntimeError(
+                    "The torchcodec backend requires TorchCodec; install it "
+                    "with `python -m pip install torchcodec`"
+                ) from exc
+            self._torchcodec = torchcodec
+            self._torchcodec_decoder_cls = VideoDecoder
         else:
             try:
                 import PyNvVideoCodec as nvc
@@ -330,6 +343,9 @@ class CausalSourceAV:
         self._decord_readers: OrderedDict[int, Any] = OrderedDict()
         self._decord_reader_opens = 0
         self._decord_reader_evictions = 0
+        self._torchcodec_readers: OrderedDict[int, Any] = OrderedDict()
+        self._torchcodec_reader_opens = 0
+        self._torchcodec_reader_evictions = 0
         self._nvc_decoder = None
         self._nvc_source_path: Path | None = None
         self._nvc_decoder_reconfigurations = 0
@@ -510,6 +526,32 @@ class CausalSourceAV:
         if self.decoder_backend == "decord":
             reader = self._decord_reader_for_clip(clip_index)
             array = reader[frame_index].asnumpy()
+        elif self.decoder_backend == "torchcodec":
+            decoder = self._torchcodec_reader_for_clip(clip_index)
+            decoded = decoder.get_frames_at(indices=[frame_index])
+            frame_tensor = decoded.data if hasattr(decoded, "data") else decoded
+            if frame_tensor.ndim == 4 and frame_tensor.shape[0] == 1:
+                frame_tensor = frame_tensor[0]
+            if frame_tensor.ndim != 3:
+                raise RuntimeError(
+                    "TorchCodec output must be one CHW or HWC frame, got "
+                    f"{tuple(frame_tensor.shape)} for {clip.path}"
+                )
+            if frame_tensor.shape[0] in (3, 4):
+                frame_tensor = frame_tensor[:3].permute(1, 2, 0)
+            elif frame_tensor.shape[-1] == 4:
+                frame_tensor = frame_tensor[..., :3]
+            elif frame_tensor.shape[-1] != 3:
+                raise RuntimeError(
+                    "TorchCodec output has no RGB channel axis: "
+                    f"{tuple(frame_tensor.shape)} for {clip.path}"
+                )
+            if frame_tensor.dtype != torch.uint8:
+                raise RuntimeError(
+                    "TorchCodec RGB output must be uint8, got "
+                    f"{frame_tensor.dtype} for {clip.path}"
+                )
+            array = frame_tensor.cpu().contiguous().numpy()
         else:
             decoder = self._pynv_decoder_for_clip(clip)
             decoded_frame = decoder[frame_index]
@@ -565,6 +607,37 @@ class CausalSourceAV:
             self._decord_reader_evictions += 1
         return reader
 
+    def _torchcodec_reader_for_clip(self, clip_index: int) -> Any:
+        """Return an LRU-managed, exact-seek TorchCodec CPU decoder."""
+
+        reader = self._torchcodec_readers.get(clip_index)
+        if reader is not None:
+            self._torchcodec_readers.move_to_end(clip_index)
+            return reader
+        if self._torchcodec_decoder_cls is None:
+            raise RuntimeError("TorchCodec was not initialized")
+
+        clip = self.clips[clip_index]
+        reader = self._torchcodec_decoder_cls(
+            str(clip.path),
+            num_ffmpeg_threads=self.decoder_threads,
+        )
+        decoded_length = getattr(reader.metadata, "num_frames", None)
+        if decoded_length is not None and int(decoded_length) != clip.num_frames:
+            raise RuntimeError(
+                "TorchCodec/ffprobe frame-count mismatch for "
+                f"{clip.path}: decoder={decoded_length}, ffprobe={clip.num_frames}"
+            )
+        self._torchcodec_readers[clip_index] = reader
+        self._torchcodec_readers.move_to_end(clip_index)
+        self._torchcodec_reader_opens += 1
+
+        while len(self._torchcodec_readers) > self.max_open_readers:
+            _, evicted_reader = self._torchcodec_readers.popitem(last=False)
+            del evicted_reader
+            self._torchcodec_reader_evictions += 1
+        return reader
+
     def _pynv_decoder_for_clip(self, clip: IndexedSourceClip) -> Any:
         """Return one reusable NVDEC decoder configured for ``clip``."""
 
@@ -652,26 +725,37 @@ class CausalSourceAV:
         }
 
     def metadata(self) -> dict[str, Any]:
-        return {
-            "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_names = {
+            "decord": "decord_with_frame_pts_and_ffmpeg_pcm_audio",
+            "torchcodec": "torchcodec_cpu_with_ffprobe_pts_and_ffmpeg_pcm_audio",
+            "pynvcodec": (
+                "pynvvideocodec_nvdec_with_ffprobe_pts_and_ffmpeg_pcm_audio"
             ),
+        }
+        return {
+            "decoder": decoder_names[self.decoder_backend],
             "decoder_backend": self.decoder_backend,
             "decoder_gpu_id": (
                 self.decoder_gpu_id if self.decoder_backend == "pynvcodec" else None
             ),
             "decoder_threads": (
-                self.decoder_threads if self.decoder_backend == "decord" else None
+                self.decoder_threads
+                if self.decoder_backend in {"decord", "torchcodec"}
+                else None
             ),
             "max_open_readers": (
-                self.max_open_readers if self.decoder_backend == "decord" else None
+                self.max_open_readers
+                if self.decoder_backend in {"decord", "torchcodec"}
+                else None
             ),
             "decoder_version": (
                 getattr(self._nvc, "__version__", None)
                 if self.decoder_backend == "pynvcodec"
+                else (
+                    getattr(self._torchcodec, "__version__", None)
+                    if self.decoder_backend == "torchcodec"
                     else None
+                )
             ),
             "visual_sampling": "global_grid_latest_source_pts_not_after_tick",
             "interval_convention": "half_open_[start,end)",
@@ -707,6 +791,8 @@ class CausalSourceAV:
                 "decoder_reconfigurations": self._nvc_decoder_reconfigurations,
                 "decord_reader_opens": self._decord_reader_opens,
                 "decord_reader_evictions": self._decord_reader_evictions,
+                "torchcodec_reader_opens": self._torchcodec_reader_opens,
+                "torchcodec_reader_evictions": self._torchcodec_reader_evictions,
             },
         }
 
@@ -835,9 +921,9 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
     parser.add_argument("--frame-cache-size", type=int, default=8)
     parser.add_argument(
         "--decoder",
-        choices=("decord", "pynvcodec"),
+        choices=("decord", "torchcodec", "pynvcodec"),
         default="decord",
-        help="Video pixel decoder; pynvcodec uses NVIDIA NVDEC",
+        help="Video pixel decoder; torchcodec is CPU and pynvcodec uses NVDEC",
     )
     parser.add_argument(
         "--decoder-gpu-id",
@@ -849,13 +935,13 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
         "--decoder-threads",
         type=int,
         default=1,
-        help="CPU decode threads per Decord VideoReader",
+        help="CPU decode threads per Decord or TorchCodec reader",
     )
     parser.add_argument(
         "--max-open-readers",
         type=int,
         default=2,
-        help="Maximum simultaneously open Decord VideoReaders per process",
+        help="Maximum simultaneously open Decord or TorchCodec readers",
     )
     parser.add_argument("--verify-every", type=int, default=0)
     parser.add_argument("--verify-continuous-every", type=int, default=0)

train_omni_turn_lora.py

@@ -113,7 +113,11 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
         help="Maximum producer-ready CPU batches queued ahead of training",
     )
     parser.add_argument("--frame-cache-size", type=int, default=64)
-    parser.add_argument("--decoder", choices=("decord", "pynvcodec"), default="decord")
+    parser.add_argument(
+        "--decoder",
+        choices=("decord", "torchcodec", "pynvcodec"),
+        default="decord",
+    )
     parser.add_argument("--decoder-gpu-id", type=int, default=0)
     parser.add_argument("--decoder-threads", type=int, default=1)
     parser.add_argument("--max-open-readers", type=int, default=1)