From c8446739845d5852ae2792de378c2281cafcc78f Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Sun, 23 Aug 2026 16:41:05 +0800 Subject: [PATCH] Add TorchCodec video backend --- cache_omni_av_pieces.py | 6 +- extract_ego4d_batch.py | 2 +- requirements.txt | 1 + streaming_turn_model_directly.py | 114 +++++++++++++++++++++++++++---- train_omni_turn_lora.py | 6 +- 5 files changed, 112 insertions(+), 17 deletions(-) diff --git a/cache_omni_av_pieces.py b/cache_omni_av_pieces.py index 1f4c130..4382e87 100644 --- a/cache_omni_av_pieces.py +++ b/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) diff --git a/extract_ego4d_batch.py b/extract_ego4d_batch.py index cffc5fb..9309129 100644 --- a/extract_ego4d_batch.py +++ b/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) diff --git a/requirements.txt b/requirements.txt index b6ea753..510e155 100644 --- a/requirements.txt +++ b/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 diff --git a/streaming_turn_model_directly.py b/streaming_turn_model_directly.py index 843c2f9..e0cbdd4 100644 --- a/streaming_turn_model_directly.py +++ b/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 None + 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) diff --git a/train_omni_turn_lora.py b/train_omni_turn_lora.py index 1409cc3..91727be 100644 --- a/train_omni_turn_lora.py +++ b/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)