torchcodec 0.3.0__cp312-cp312-macosx_11_0_arm64.whl

torchcodec/_core/fetch_and_expose_non_gpl_ffmpeg_libs.cmake

@@ -0,0 +1,226 @@
+# This file fetches the non-GPL ffmpeg libraries from the torchcodec S3 bucket,
+# and exposes them as CMake targets so we can dynamically link against them.
+# These libraries were built on the CI via the build_ffmpeg.yaml workflow.
+
+# Avoid warning: see https://cmake.org/cmake/help/latest/policy/CMP0135.html
+if (CMAKE_VERSION VERSION_GREATER_EQUAL "3.24.0")
+    cmake_policy(SET CMP0135 NEW)
+endif()
+
+include(FetchContent)
+
+set(
+    base_url
+    https://pytorch.s3.amazonaws.com/torchcodec/ffmpeg/2025-03-14
+)
+
+if (${CMAKE_SYSTEM_NAME} STREQUAL "Linux")
+    set(
+        platform_url
+        ${base_url}/linux_x86_64
+    )
+
+    set(
+        f4_sha256
+        1a083f1922443bedb5243d04896383b8c606778a7ddb9d886c8303e55339fe0c
+    )
+    set(
+        f5_sha256
+        65d6ad54082d94dcb3f801d73df2265e0e1bb303c7afbce7723e3b77ccd0e207
+    )
+    set(
+        f6_sha256
+        8bd5939c2f4a4b072e837e7870c13fe7d13824e5ff087ab534e4db4e90b7be9c
+    )
+    set(
+        f7_sha256
+        1cb946d8b7c6393c2c3ebe1f900b8de7a2885fe614c45d4ec32c9833084f2f26
+    )
+
+    set(
+       f4_library_file_names
+       libavutil.so.56
+       libavcodec.so.58
+       libavformat.so.58
+       libavdevice.so.58
+       libavfilter.so.7
+       libswscale.so.5
+       libswresample.so.3
+    )
+    set(
+       f5_library_file_names
+       libavutil.so.57
+       libavcodec.so.59
+       libavformat.so.59
+       libavdevice.so.59
+       libavfilter.so.8
+       libswscale.so.6
+       libswresample.so.4
+    )
+    set(
+       f6_library_file_names
+       libavutil.so.58
+       libavcodec.so.60
+       libavformat.so.60
+       libavdevice.so.60
+       libavfilter.so.9
+       libswscale.so.7
+       libswresample.so.4
+    )
+    set(
+       f7_library_file_names
+       libavutil.so.59
+       libavcodec.so.61
+       libavformat.so.61
+       libavdevice.so.61
+       libavfilter.so.10
+       libswscale.so.8
+       libswresample.so.5
+    )
+elseif (${CMAKE_SYSTEM_NAME} STREQUAL "Darwin")
+    set(
+        platform_url
+        ${base_url}/macos_arm64
+    )
+    set(
+        f4_sha256
+        f0335434529d9e19359eae0fe912dd9e747667534a1c92e662f5219a55dfad8c
+    )
+    set(
+        f5_sha256
+        cfc3449c9af6863731a431ce89e32c08c5f8ece94b306fb6b695828502a76166
+    )
+    set(
+        f6_sha256
+        ec47b4783c342038e720e33b2fdfa55a9a490afb1cf37a26467733983688647e
+    )
+    set(
+        f7_sha256
+        48a4fc8ce098305cfd4a58f40889249c523ca3c285f66ba704b5bad0e3ada53a
+    )
+    set(
+       f4_library_file_names
+       libavutil.56.dylib
+       libavcodec.58.dylib
+       libavformat.58.dylib
+       libavdevice.58.dylib
+       libavfilter.7.dylib
+       libswscale.5.dylib
+       libswresample.3.dylib
+    )
+    set(
+       f5_library_file_names
+       libavutil.57.dylib
+       libavcodec.59.dylib
+       libavformat.59.dylib
+       libavdevice.59.dylib
+       libavfilter.8.dylib
+       libswscale.6.dylib
+       libswresample.4.dylib
+    )
+    set(
+       f6_library_file_names
+       libavutil.58.dylib
+       libavcodec.60.dylib
+       libavformat.60.dylib
+       libavdevice.60.dylib
+       libavfilter.9.dylib
+       libswscale.7.dylib
+       libswresample.4.dylib
+    )
+    set(
+       f7_library_file_names
+       libavutil.59.dylib
+       libavcodec.61.dylib
+       libavformat.61.dylib
+       libavdevice.61.dylib
+       libavfilter.10.dylib
+       libswscale.8.dylib
+       libswresample.5.dylib
+    )
+else()
+    message(
+        FATAL_ERROR
+        "Unsupported operating system: ${CMAKE_SYSTEM_NAME}"
+    )
+endif()
+
+FetchContent_Declare(
+    f4
+    URL ${platform_url}/4.4.4.tar.gz
+    URL_HASH
+    SHA256=${f4_sha256}
+)
+FetchContent_Declare(
+    f5
+    URL ${platform_url}/5.1.4.tar.gz
+    URL_HASH
+    SHA256=${f5_sha256}
+)
+FetchContent_Declare(
+    f6
+    URL ${platform_url}/6.1.1.tar.gz
+    URL_HASH
+    SHA256=${f6_sha256}
+)
+FetchContent_Declare(
+    f7
+    URL ${platform_url}/7.0.1.tar.gz
+    URL_HASH
+    SHA256=${f7_sha256}
+)
+
+FetchContent_MakeAvailable(f4 f5 f6 f7)
+
+add_library(ffmpeg4 INTERFACE)
+add_library(ffmpeg5 INTERFACE)
+add_library(ffmpeg6 INTERFACE)
+add_library(ffmpeg7 INTERFACE)
+
+# Note: the f?_SOURCE_DIR variables were set by FetchContent_MakeAvailable
+target_include_directories(ffmpeg4 INTERFACE ${f4_SOURCE_DIR}/include)
+target_include_directories(ffmpeg5 INTERFACE ${f5_SOURCE_DIR}/include)
+target_include_directories(ffmpeg6 INTERFACE ${f6_SOURCE_DIR}/include)
+target_include_directories(ffmpeg7 INTERFACE ${f7_SOURCE_DIR}/include)
+
+list(
+    TRANSFORM f4_library_file_names
+    PREPEND ${f4_SOURCE_DIR}/lib/
+    OUTPUT_VARIABLE f4_library_paths
+)
+list(
+    TRANSFORM f5_library_file_names
+    PREPEND ${f5_SOURCE_DIR}/lib/
+    OUTPUT_VARIABLE f5_library_paths
+)
+list(
+    TRANSFORM f6_library_file_names
+    PREPEND ${f6_SOURCE_DIR}/lib/
+    OUTPUT_VARIABLE f6_library_paths
+)
+list(
+    TRANSFORM f7_library_file_names
+    PREPEND ${f7_SOURCE_DIR}/lib/
+    OUTPUT_VARIABLE f7_library_paths
+)
+
+target_link_libraries(
+    ffmpeg4
+    INTERFACE
+    ${f4_library_paths}
+)
+target_link_libraries(
+    ffmpeg5
+    INTERFACE
+    ${f5_library_paths}
+)
+target_link_libraries(
+    ffmpeg6
+    INTERFACE
+    ${f6_library_paths}
+)
+target_link_libraries(
+    ffmpeg7
+    INTERFACE
+    ${f7_library_paths}
+)

torchcodec/_core/ops.py

@@ -0,0 +1,381 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import io
+import json
+import warnings
+from types import ModuleType
+from typing import List, Optional, Tuple, Union
+
+import torch
+from torch.library import get_ctx, register_fake
+
+from torchcodec._internally_replaced_utils import (  # @manual=//pytorch/torchcodec/src:internally_replaced_utils
+    _get_extension_path,
+    _load_pybind11_module,
+)
+
+_pybind_ops: Optional[ModuleType] = None
+
+
+def load_torchcodec_shared_libraries():
+    # Successively try to load libtorchcodec_*7.so, libtorchcodec_*6.so,
+    # libtorchcodec_*5.so, and libtorchcodec_*4.so. Each of these correspond to an
+    # ffmpeg major version. This should cover all potential ffmpeg versions
+    # installed on the user's machine.
+    #
+    # On fbcode, _get_extension_path() is overridden and directly points to the
+    # correct .so file, so this for-loop succeeds on the first iteration.
+    #
+    # Note that we use two different methods for loading shared libraries:
+    #
+    #   1. torch.ops.load_library(): For PyTorch custom ops and the C++ only
+    #      libraries the custom ops depend on. Loading libraries through PyTorch
+    #      registers the custom ops with PyTorch's runtime and the ops can be
+    #      accessed through torch.ops after loading.
+    #
+    #   2. importlib: For pybind11 modules. We load them dynamically, rather
+    #      than using a plain import statement. A plain import statement only
+    #      works when the module name and file name match exactly. Our shared
+    #      libraries do not meet those conditions.
+
+    exceptions = []
+    pybind_ops_module_name = "decoder_core_pybind_ops"
+    for ffmpeg_major_version in (7, 6, 5, 4):
+        decoder_library_name = f"libtorchcodec_decoder{ffmpeg_major_version}"
+        custom_ops_library_name = f"libtorchcodec_custom_ops{ffmpeg_major_version}"
+        pybind_ops_library_name = f"libtorchcodec_pybind_ops{ffmpeg_major_version}"
+        try:
+            torch.ops.load_library(_get_extension_path(decoder_library_name))
+            torch.ops.load_library(_get_extension_path(custom_ops_library_name))
+
+            pybind_ops_library_path = _get_extension_path(pybind_ops_library_name)
+            global _pybind_ops
+            _pybind_ops = _load_pybind11_module(
+                pybind_ops_module_name, pybind_ops_library_path
+            )
+            return
+        except Exception as e:
+            # TODO: recording and reporting exceptions this way is OK for now as  it's just for debugging,
+            # but we should probably handle that via a proper logging mechanism.
+            exceptions.append((ffmpeg_major_version, e))
+
+    traceback = (
+        "\n[start of libtorchcodec loading traceback]\n"
+        + "\n".join(f"FFmpeg version {v}: {str(e)}" for v, e in exceptions)
+        + "\n[end of libtorchcodec loading traceback]."
+    )
+    raise RuntimeError(
+        f"""Could not load libtorchcodec. Likely causes:
+          1. FFmpeg is not properly installed in your environment. We support
+             versions 4, 5, 6 and 7.
+          2. The PyTorch version ({torch.__version__}) is not compatible with
+             this version of TorchCodec. Refer to the version compatibility
+             table:
+             https://github.com/pytorch/torchcodec?tab=readme-ov-file#installing-torchcodec.
+          3. Another runtime dependency; see exceptions below.
+        The following exceptions were raised as we tried to load libtorchcodec:
+        """
+        f"{traceback}"
+    )
+
+
+load_torchcodec_shared_libraries()
+
+
+# Note: We use disallow_in_graph because PyTorch does constant propagation of
+# factory functions.
+create_from_file = torch._dynamo.disallow_in_graph(
+    torch.ops.torchcodec_ns.create_from_file.default
+)
+create_audio_encoder = torch._dynamo.disallow_in_graph(
+    torch.ops.torchcodec_ns.create_audio_encoder.default
+)
+encode_audio = torch._dynamo.disallow_in_graph(
+    torch.ops.torchcodec_ns.encode_audio.default
+)
+create_from_tensor = torch._dynamo.disallow_in_graph(
+    torch.ops.torchcodec_ns.create_from_tensor.default
+)
+_convert_to_tensor = torch._dynamo.disallow_in_graph(
+    torch.ops.torchcodec_ns._convert_to_tensor.default
+)
+add_video_stream = torch.ops.torchcodec_ns.add_video_stream.default
+_add_video_stream = torch.ops.torchcodec_ns._add_video_stream.default
+add_audio_stream = torch.ops.torchcodec_ns.add_audio_stream.default
+seek_to_pts = torch.ops.torchcodec_ns.seek_to_pts.default
+get_next_frame = torch.ops.torchcodec_ns.get_next_frame.default
+get_frame_at_pts = torch.ops.torchcodec_ns.get_frame_at_pts.default
+get_frame_at_index = torch.ops.torchcodec_ns.get_frame_at_index.default
+get_frames_at_indices = torch.ops.torchcodec_ns.get_frames_at_indices.default
+get_frames_by_pts = torch.ops.torchcodec_ns.get_frames_by_pts.default
+get_frames_in_range = torch.ops.torchcodec_ns.get_frames_in_range.default
+get_frames_by_pts_in_range = torch.ops.torchcodec_ns.get_frames_by_pts_in_range.default
+get_frames_by_pts_in_range_audio = (
+    torch.ops.torchcodec_ns.get_frames_by_pts_in_range_audio.default
+)
+get_json_metadata = torch.ops.torchcodec_ns.get_json_metadata.default
+_test_frame_pts_equality = torch.ops.torchcodec_ns._test_frame_pts_equality.default
+_get_container_json_metadata = (
+    torch.ops.torchcodec_ns.get_container_json_metadata.default
+)
+_get_key_frame_indices = torch.ops.torchcodec_ns._get_key_frame_indices.default
+scan_all_streams_to_update_metadata = (
+    torch.ops.torchcodec_ns.scan_all_streams_to_update_metadata.default
+)
+_get_stream_json_metadata = torch.ops.torchcodec_ns.get_stream_json_metadata.default
+_get_json_ffmpeg_library_versions = (
+    torch.ops.torchcodec_ns._get_json_ffmpeg_library_versions.default
+)
+
+
+# =============================
+# Functions not related to custom ops, but similar implementation to c++ ops
+# =============================
+def create_from_bytes(
+    video_bytes: bytes, seek_mode: Optional[str] = None
+) -> torch.Tensor:
+    with warnings.catch_warnings():
+        # Ignore warning stating that the underlying video_bytes buffer is
+        # non-writable.
+        warnings.filterwarnings("ignore", category=UserWarning)
+        buffer = torch.frombuffer(video_bytes, dtype=torch.uint8)
+    return create_from_tensor(buffer, seek_mode)
+
+
+def create_from_file_like(
+    file_like: Union[io.RawIOBase, io.BufferedReader], seek_mode: Optional[str] = None
+) -> torch.Tensor:
+    assert _pybind_ops is not None
+    return _convert_to_tensor(_pybind_ops.create_from_file_like(file_like, seek_mode))
+
+
+# ==============================
+# Abstract impl for the operators. Needed by torch.compile.
+# ==============================
+@register_fake("torchcodec_ns::create_from_file")
+def create_from_file_abstract(filename: str, seek_mode: Optional[str]) -> torch.Tensor:
+    return torch.empty([], dtype=torch.long)
+
+
+@register_fake("torchcodec_ns::create_audio_encoder")
+def create_audio_encoder_abstract(
+    wf: torch.Tensor, sample_rate: int, filename: str, bit_rate: Optional[int] = None
+) -> torch.Tensor:
+    return torch.empty([], dtype=torch.long)
+
+
+@register_fake("torchcodec_ns::encode_audio")
+def encode_audio_abstract(encoder: torch.Tensor) -> torch.Tensor:
+    return torch.empty([], dtype=torch.long)
+
+
+@register_fake("torchcodec_ns::create_from_tensor")
+def create_from_tensor_abstract(
+    video_tensor: torch.Tensor, seek_mode: Optional[str]
+) -> torch.Tensor:
+    return torch.empty([], dtype=torch.long)
+
+
+@register_fake("torchcodec_ns::_convert_to_tensor")
+def _convert_to_tensor_abstract(decoder_ptr: int) -> torch.Tensor:
+    return torch.empty([], dtype=torch.long)
+
+
+@register_fake("torchcodec_ns::_add_video_stream")
+def _add_video_stream_abstract(
+    decoder: torch.Tensor,
+    *,
+    width: Optional[int] = None,
+    height: Optional[int] = None,
+    num_threads: Optional[int] = None,
+    dimension_order: Optional[str] = None,
+    stream_index: Optional[int] = None,
+    device: Optional[str] = None,
+    color_conversion_library: Optional[str] = None,
+) -> None:
+    return
+
+
+@register_fake("torchcodec_ns::add_video_stream")
+def add_video_stream_abstract(
+    decoder: torch.Tensor,
+    *,
+    width: Optional[int] = None,
+    height: Optional[int] = None,
+    num_threads: Optional[int] = None,
+    dimension_order: Optional[str] = None,
+    stream_index: Optional[int] = None,
+    device: Optional[str] = None,
+) -> None:
+    return
+
+
+@register_fake("torchcodec_ns::add_audio_stream")
+def add_audio_stream_abstract(
+    decoder: torch.Tensor,
+    *,
+    stream_index: Optional[int] = None,
+) -> None:
+    return
+
+
+@register_fake("torchcodec_ns::seek_to_pts")
+def seek_abstract(decoder: torch.Tensor, seconds: float) -> None:
+    return
+
+
+@register_fake("torchcodec_ns::get_next_frame")
+def get_next_frame_abstract(
+    decoder: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    # Images are 3 dimensions: height, width, channels.
+    # The exact permutation depends on the constructor options passed in.
+    image_size = [get_ctx().new_dynamic_size() for _ in range(3)]
+    return (
+        torch.empty(image_size),
+        torch.empty([], dtype=torch.float),
+        torch.empty([], dtype=torch.float),
+    )
+
+
+@register_fake("torchcodec_ns::get_frame_at_pts")
+def get_frame_at_pts_abstract(
+    decoder: torch.Tensor, seconds: float
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    image_size = [get_ctx().new_dynamic_size() for _ in range(3)]
+    return (
+        torch.empty(image_size),
+        torch.empty([], dtype=torch.float),
+        torch.empty([], dtype=torch.float),
+    )
+
+
+@register_fake("torchcodec_ns::get_frames_by_pts")
+def get_frames_by_pts_abstract(
+    decoder: torch.Tensor,
+    *,
+    timestamps: List[float],
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    image_size = [get_ctx().new_dynamic_size() for _ in range(4)]
+    return (
+        torch.empty(image_size),
+        torch.empty([], dtype=torch.float),
+        torch.empty([], dtype=torch.float),
+    )
+
+
+@register_fake("torchcodec_ns::get_frame_at_index")
+def get_frame_at_index_abstract(
+    decoder: torch.Tensor, *, frame_index: int
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    image_size = [get_ctx().new_dynamic_size() for _ in range(3)]
+    return (
+        torch.empty(image_size),
+        torch.empty([], dtype=torch.float),
+        torch.empty([], dtype=torch.float),
+    )
+
+
+@register_fake("torchcodec_ns::get_frames_at_indices")
+def get_frames_at_indices_abstract(
+    decoder: torch.Tensor,
+    *,
+    frame_indices: List[int],
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    image_size = [get_ctx().new_dynamic_size() for _ in range(4)]
+    return (
+        torch.empty(image_size),
+        torch.empty([], dtype=torch.float),
+        torch.empty([], dtype=torch.float),
+    )
+
+
+@register_fake("torchcodec_ns::get_frames_in_range")
+def get_frames_in_range_abstract(
+    decoder: torch.Tensor,
+    *,
+    start: int,
+    stop: int,
+    step: Optional[int] = None,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    image_size = [get_ctx().new_dynamic_size() for _ in range(4)]
+    return (
+        torch.empty(image_size),
+        torch.empty([], dtype=torch.float),
+        torch.empty([], dtype=torch.float),
+    )
+
+
+@register_fake("torchcodec_ns::get_frames_by_pts_in_range")
+def get_frames_by_pts_in_range_abstract(
+    decoder: torch.Tensor,
+    *,
+    start_seconds: float,
+    stop_seconds: float,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    image_size = [get_ctx().new_dynamic_size() for _ in range(4)]
+    return (
+        torch.empty(image_size),
+        torch.empty([], dtype=torch.float),
+        torch.empty([], dtype=torch.float),
+    )
+
+
+@register_fake("torchcodec_ns::get_frames_by_pts_in_range_audio")
+def get_frames_by_pts_in_range_audio_abstract(
+    decoder: torch.Tensor,
+    *,
+    start_seconds: float,
+    stop_seconds: Optional[float] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    image_size = [get_ctx().new_dynamic_size() for _ in range(4)]
+    return (torch.empty(image_size), torch.empty([], dtype=torch.float))
+
+
+@register_fake("torchcodec_ns::_get_key_frame_indices")
+def get_key_frame_indices_abstract(decoder: torch.Tensor) -> torch.Tensor:
+    return torch.empty([], dtype=torch.int)
+
+
+@register_fake("torchcodec_ns::get_json_metadata")
+def get_json_metadata_abstract(decoder: torch.Tensor) -> str:
+    return ""
+
+
+@register_fake("torchcodec_ns::get_container_json_metadata")
+def get_container_json_metadata_abstract(decoder: torch.Tensor) -> str:
+    return ""
+
+
+@register_fake("torchcodec_ns::get_stream_json_metadata")
+def get_stream_json_metadata_abstract(decoder: torch.Tensor, stream_idx: int) -> str:
+    return ""
+
+
+@register_fake("torchcodec_ns::_test_frame_pts_equality")
+def _test_frame_pts_equality_abstract(
+    decoder: torch.Tensor,
+    *,
+    frame_index: int,
+    pts_seconds_to_test: float,
+) -> bool:
+    return False
+
+
+@register_fake("torchcodec_ns::_get_json_ffmpeg_library_versions")
+def _get_json_ffmpeg_library_versions_abstract() -> str:
+    return ""
+
+
+@register_fake("torchcodec_ns::scan_all_streams_to_update_metadata")
+def scan_all_streams_to_update_metadata_abstract(decoder: torch.Tensor) -> None:
+    return
+
+
+def get_ffmpeg_library_versions():
+    versions_json = _get_json_ffmpeg_library_versions()
+    return json.loads(versions_json)

torchcodec/_core/pybind_ops.cpp

@@ -0,0 +1,45 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <cstdint>
+#include <string>
+
+#include "src/torchcodec/_core/AVIOFileLikeContext.h"
+#include "src/torchcodec/_core/SingleStreamDecoder.h"
+
+namespace py = pybind11;
+
+namespace facebook::torchcodec {
+
+// In principle, this should be able to return a tensor. But when we try that,
+// we run into the bug reported here:
+//
+//   https://github.com/pytorch/pytorch/issues/136664
+//
+// So we instead launder the pointer through an int, and then use a conversion
+// function on the custom ops side to launder that int into a tensor.
+int64_t create_from_file_like(
+    py::object file_like,
+    std::optional<std::string_view> seek_mode) {
+  SingleStreamDecoder::SeekMode realSeek = SingleStreamDecoder::SeekMode::exact;
+  if (seek_mode.has_value()) {
+    realSeek = seekModeFromString(seek_mode.value());
+  }
+
+  auto avioContextHolder = std::make_unique<AVIOFileLikeContext>(file_like);
+
+  SingleStreamDecoder* decoder =
+      new SingleStreamDecoder(std::move(avioContextHolder), realSeek);
+  return reinterpret_cast<int64_t>(decoder);
+}
+
+PYBIND11_MODULE(decoder_core_pybind_ops, m) {
+  m.def("create_from_file_like", &create_from_file_like);
+}
+
+} // namespace facebook::torchcodec