diff --git a/README.md b/README.md index c122bb3..3f42b01 100644 --- a/README.md +++ b/README.md @@ -651,6 +651,35 @@ python streaming_turn_model_directly.py \ --output features/direct/Session_1.pt ``` +The direct reader defaults to Decord. On NVIDIA systems, install +`PyNvVideoCodec` and select the NVDEC backend explicitly: + +```bash +python -m pip install PyNvVideoCodec + +CUDA_VISIBLE_DEVICES=0 python streaming_turn_model_directly.py \ + --model-id /path/to/Qwen2.5-Omni-3B \ + --easycom-session-dir /path/to/Main/Video_Compressed/Session_1 \ + --end 24 \ + --mode open-stream-replay \ + --memory-seconds 20 \ + --video-fps 2 \ + --frame-cache-size 64 \ + --decoder pynvcodec \ + --decoder-gpu-id 0 \ + --flash-attention-2 \ + --device cuda:0 \ + --output features/direct/Session_1_24s_pynvcodec.pt +``` + +The NVDEC backend uses ffprobe for exact display-order PTS indexing, one +reconfigurable `SimpleDecoder` for pixel decode, and the existing ffmpeg PCM +audio path. `--decoder-gpu-id` is relative to `CUDA_VISIBLE_DEVICES`, so it is +normally `0` even when the worker is assigned a nonzero host GPU. Decoder RGB +conversion can differ slightly; validate a short Decord/PyNvVideoCodec pair +before extracting every session, and do not mix decoder backends across +train/test splits. + Matching `states.shape[0]` to the target is necessary but not sufficient. The direct reader also validates EasyCom's 20-Hz source inventory, source PTS monotonicity, dense duration, crop bounds, audio sample count, and causal diff --git a/requirements.txt b/requirements.txt index 0690fb7..2419e09 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,3 +9,4 @@ tqdm scikit-learn # Optional on compatible NVIDIA GPUs: flash-attn +# Optional NVDEC backend: PyNvVideoCodec diff --git a/run_extract_easycom_2x4090_decord.sh b/run_extract_easycom_2x4090_decord.sh new file mode 100644 index 0000000..4648028 --- /dev/null +++ b/run_extract_easycom_2x4090_decord.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +# Extract all 12 EasyCom sessions on host GPUs 3 and 4. +# Exactly one Qwen/Decord worker runs on each GPU; sessions are processed +# sequentially within each worker. The ordering balances the official 5-Hz +# target lengths across the two GPUs. +set -euo pipefail + +PROJECT_ROOT="${PROJECT_ROOT:-/data/wangx/Code/icassp2027-experiment}" +DATASET_ROOT="${DATASET_ROOT:-/data/wangx/Datasets/easycom}" +MODEL="${MODEL:-/data/wangx/Models/Qwen2.5-Omni-3B}" + +VIDEO_ROOT="${VIDEO_ROOT:-${DATASET_ROOT}/extracted_easycom/Main/Video_Compressed}" +TARGET_ROOT="${TARGET_ROOT:-${DATASET_ROOT}/easycom_target_frame}" +FEATURE_ROOT="${FEATURE_ROOT:-${DATASET_ROOT}/omni_features/open_stream_2fps_20s_decord}" +LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/easycom_2x4090_decord}" + +mkdir -p "${FEATURE_ROOT}" "${LOG_ROOT}" + +for gpu in 3 4; do + nvidia-smi -i "${gpu}" --query-gpu=index,name,memory.total \ + --format=csv,noheader +done + +for sid in {1..12}; do + test -d "${VIDEO_ROOT}/Session_${sid}" || { + echo "Missing video directory: ${VIDEO_ROOT}/Session_${sid}" >&2 + exit 1 + } + test -s "${TARGET_ROOT}/Session_${sid}.npy" || { + echo "Missing target: ${TARGET_ROOT}/Session_${sid}.npy" >&2 + exit 1 + } +done + +echo "Features: ${FEATURE_ROOT}" +echo "Logs: ${LOG_ROOT}" +echo "GPU 3: Session_12 Session_5 Session_9 Session_8 Session_10 Session_1" +echo "GPU 4: Session_6 Session_4 Session_3 Session_2 Session_7 Session_11" + +exec nice -n 5 env \ + OMP_NUM_THREADS=4 \ + MKL_NUM_THREADS=4 \ + OPENBLAS_NUM_THREADS=4 \ + PROJECT_ROOT="${PROJECT_ROOT}" \ + MODEL="${MODEL}" \ + VIDEO_ROOT="${VIDEO_ROOT}" \ + TARGET_ROOT="${TARGET_ROOT}" \ + FEATURE_ROOT="${FEATURE_ROOT}" \ + LOG_ROOT="${LOG_ROOT}" \ + GPUS="3 4" \ + SESSIONS="12 6 5 4 9 3 8 2 10 7 1 11" \ + MEMORY_SECONDS=20 \ + NATIVE_BLOCK_SECONDS=2 \ + VIDEO_FPS=2 \ + FRAME_CACHE_SIZE=64 \ + MODE=open-stream-replay \ + DTYPE=bfloat16 \ + DECODER=decord \ + DECODER_GPU_ID=0 \ + bash "${PROJECT_ROOT}/run_extract_easycom_direct.sh" diff --git a/run_extract_easycom_direct.sh b/run_extract_easycom_direct.sh index c95d420..50d6c20 100644 --- a/run_extract_easycom_direct.sh +++ b/run_extract_easycom_direct.sh @@ -18,6 +18,9 @@ NATIVE_BLOCK_SECONDS="${NATIVE_BLOCK_SECONDS:-2}" VIDEO_FPS="${VIDEO_FPS:-2}" MODE="${MODE:-open-stream-replay}" DTYPE="${DTYPE:-bfloat16}" +DECODER="${DECODER:-decord}" +DECODER_GPU_ID="${DECODER_GPU_ID:-0}" +FRAME_CACHE_SIZE="${FRAME_CACHE_SIZE:-64}" # Space-separated host GPU ids, e.g. "0" or "0 1 2 3". One session per GPU at a time. GPUS="${GPUS:-0}" # Optional: only run selected session ids, e.g. SESSIONS="1 2 3" @@ -69,6 +72,7 @@ run_one() { echo "target=${target}" echo "output=${output}" echo "memory_seconds=${MEMORY_SECONDS} video_fps=${VIDEO_FPS} mode=${MODE}" + echo "decoder=${DECODER} decoder_gpu_id=${DECODER_GPU_ID} frame_cache_size=${FRAME_CACHE_SIZE}" if ! CUDA_VISIBLE_DEVICES="${gpu}" PYTHONUNBUFFERED=1 \ python streaming_turn_model_directly.py \ @@ -80,6 +84,9 @@ run_one() { --memory-seconds "${MEMORY_SECONDS}" \ --native-block-seconds "${NATIVE_BLOCK_SECONDS}" \ --video-fps "${VIDEO_FPS}" \ + --frame-cache-size "${FRAME_CACHE_SIZE}" \ + --decoder "${DECODER}" \ + --decoder-gpu-id "${DECODER_GPU_ID}" \ --dtype "${DTYPE}" \ --flash-attention-2 \ --device cuda:0 \ diff --git a/streaming_turn_model_directly.py b/streaming_turn_model_directly.py index 54e6351..c108417 100644 --- a/streaming_turn_model_directly.py +++ b/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