torchcodec 0.3.0__cp313-cp313t-manylinux_2_28_x86_64.whl

torchcodec/__init__.py

@@ -0,0 +1,16 @@
+# 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.
+
+# Note: usort wants to put Frame and FrameBatch after decoders and samplers,
+# but that results in circular import.
+from ._frame import AudioSamples, Frame, FrameBatch  # usort:skip # noqa
+from . import decoders, samplers  # noqa
+
+try:
+    # Note that version.py is generated during install.
+    from .version import __version__  # noqa: F401
+except Exception:
+    pass

torchcodec/_core/AVIOBytesContext.cpp

@@ -0,0 +1,70 @@
+// 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 "src/torchcodec/_core/AVIOBytesContext.h"
+#include <torch/types.h>
+
+namespace facebook::torchcodec {
+
+AVIOBytesContext::AVIOBytesContext(const void* data, int64_t dataSize)
+    : dataContext_{static_cast<const uint8_t*>(data), dataSize, 0} {
+  TORCH_CHECK(data != nullptr, "Video data buffer cannot be nullptr!");
+  TORCH_CHECK(dataSize > 0, "Video data size must be positive");
+  createAVIOContext(&read, &seek, &dataContext_);
+}
+
+// The signature of this function is defined by FFMPEG.
+int AVIOBytesContext::read(void* opaque, uint8_t* buf, int buf_size) {
+  auto dataContext = static_cast<DataContext*>(opaque);
+  TORCH_CHECK(
+      dataContext->current <= dataContext->size,
+      "Tried to read outside of the buffer: current=",
+      dataContext->current,
+      ", size=",
+      dataContext->size);
+
+  int64_t numBytesRead = std::min(
+      static_cast<int64_t>(buf_size), dataContext->size - dataContext->current);
+
+  TORCH_CHECK(
+      numBytesRead >= 0,
+      "Tried to read negative bytes: numBytesRead=",
+      numBytesRead,
+      ", size=",
+      dataContext->size,
+      ", current=",
+      dataContext->current);
+
+  if (numBytesRead == 0) {
+    return AVERROR_EOF;
+  }
+
+  std::memcpy(buf, dataContext->data + dataContext->current, numBytesRead);
+  dataContext->current += numBytesRead;
+  return numBytesRead;
+}
+
+// The signature of this function is defined by FFMPEG.
+int64_t AVIOBytesContext::seek(void* opaque, int64_t offset, int whence) {
+  auto dataContext = static_cast<DataContext*>(opaque);
+  int64_t ret = -1;
+
+  switch (whence) {
+    case AVSEEK_SIZE:
+      ret = dataContext->size;
+      break;
+    case SEEK_SET:
+      dataContext->current = offset;
+      ret = offset;
+      break;
+    default:
+      break;
+  }
+
+  return ret;
+}
+
+} // namespace facebook::torchcodec

torchcodec/_core/AVIOBytesContext.h

@@ -0,0 +1,32 @@
+// 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.
+
+#pragma once
+
+#include "src/torchcodec/_core/AVIOContextHolder.h"
+
+namespace facebook::torchcodec {
+
+// Enables users to pass in the entire video as bytes. Our read and seek
+// functions then traverse the bytes in memory.
+class AVIOBytesContext : public AVIOContextHolder {
+ public:
+  explicit AVIOBytesContext(const void* data, int64_t dataSize);
+
+ private:
+  struct DataContext {
+    const uint8_t* data;
+    int64_t size;
+    int64_t current;
+  };
+
+  static int read(void* opaque, uint8_t* buf, int buf_size);
+  static int64_t seek(void* opaque, int64_t offset, int whence);
+
+  DataContext dataContext_;
+};
+
+} // namespace facebook::torchcodec

torchcodec/_core/AVIOContextHolder.cpp

@@ -0,0 +1,50 @@
+// 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 "src/torchcodec/_core/AVIOContextHolder.h"
+#include <torch/types.h>
+
+namespace facebook::torchcodec {
+
+void AVIOContextHolder::createAVIOContext(
+    AVIOReadFunction read,
+    AVIOSeekFunction seek,
+    void* heldData,
+    int bufferSize) {
+  TORCH_CHECK(
+      bufferSize > 0,
+      "Buffer size must be greater than 0; is " + std::to_string(bufferSize));
+  auto buffer = static_cast<uint8_t*>(av_malloc(bufferSize));
+  TORCH_CHECK(
+      buffer != nullptr,
+      "Failed to allocate buffer of size " + std::to_string(bufferSize));
+
+  avioContext_.reset(avio_alloc_context(
+      buffer,
+      bufferSize,
+      0,
+      heldData,
+      read,
+      nullptr, // write function; not supported yet
+      seek));
+
+  if (!avioContext_) {
+    av_freep(&buffer);
+    TORCH_CHECK(false, "Failed to allocate AVIOContext");
+  }
+}
+
+AVIOContextHolder::~AVIOContextHolder() {
+  if (avioContext_) {
+    av_freep(&avioContext_->buffer);
+  }
+}
+
+AVIOContext* AVIOContextHolder::getAVIOContext() {
+  return avioContext_.get();
+}
+
+} // namespace facebook::torchcodec

torchcodec/_core/AVIOContextHolder.h

@@ -0,0 +1,65 @@
+// 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.
+
+#pragma once
+
+#include "src/torchcodec/_core/FFMPEGCommon.h"
+
+namespace facebook::torchcodec {
+
+// The AVIOContextHolder serves several purposes:
+//
+//   1. It is a smart pointer for the AVIOContext. It has the logic to create
+//      a new AVIOContext and will appropriately free the AVIOContext when it
+//      goes out of scope. Note that this requires more than just having a
+//      UniqueAVIOContext, as the AVIOContext points to a buffer which must be
+//      freed.
+//   2. It is a base class for AVIOContext specializations. When specializing a
+//      AVIOContext, we need to provide four things:
+//        1. A read callback function.
+//        2. A seek callback function.
+//        3. A write callback function. (Not supported yet; it's for encoding.)
+//        4. A pointer to some context object that has the same lifetime as the
+//           AVIOContext itself. This context object holds the custom state that
+//           tracks the custom behavior of reading, seeking and writing. It is
+//           provided upon AVIOContext creation and to the read, seek and
+//           write callback functions.
+//      While it's not required, it is natural for the derived classes to make
+//      all of the above members. Base classes need to call
+//      createAVIOContext(), ideally in their constructor.
+//  3. A generic handle for those that just need to manage having access to an
+//     AVIOContext, but aren't necessarily concerned with how it was customized:
+//     typically, the SingleStreamDecoder.
+class AVIOContextHolder {
+ public:
+  virtual ~AVIOContextHolder();
+  AVIOContext* getAVIOContext();
+
+ protected:
+  // Make constructor protected to prevent anyone from constructing
+  // an AVIOContextHolder without deriving it. (Ordinarily this would be
+  // enforced by having a pure virtual methods, but we don't have any.)
+  AVIOContextHolder() = default;
+
+  // These signatures are defined by FFmpeg.
+  using AVIOReadFunction = int (*)(void*, uint8_t*, int);
+  using AVIOSeekFunction = int64_t (*)(void*, int64_t, int);
+
+  // Deriving classes should call this function in their constructor.
+  void createAVIOContext(
+      AVIOReadFunction read,
+      AVIOSeekFunction seek,
+      void* heldData,
+      int bufferSize = defaultBufferSize);
+
+ private:
+  UniqueAVIOContext avioContext_;
+
+  // Defaults to 64 KB
+  static const int defaultBufferSize = 64 * 1024;
+};
+
+} // namespace facebook::torchcodec

torchcodec/_core/AVIOFileLikeContext.cpp

@@ -0,0 +1,80 @@
+// 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 "src/torchcodec/_core/AVIOFileLikeContext.h"
+#include <torch/types.h>
+
+namespace facebook::torchcodec {
+
+AVIOFileLikeContext::AVIOFileLikeContext(py::object fileLike)
+    : fileLike_{UniquePyObject(new py::object(fileLike))} {
+  {
+    // TODO: Is it necessary to acquire the GIL here? Is it maybe even
+    // harmful? At the moment, this is only called from within a pybind
+    // function, and pybind guarantees we have the GIL.
+    py::gil_scoped_acquire gil;
+    TORCH_CHECK(
+        py::hasattr(fileLike, "read"),
+        "File like object must implement a read method.");
+    TORCH_CHECK(
+        py::hasattr(fileLike, "seek"),
+        "File like object must implement a seek method.");
+  }
+  createAVIOContext(&read, &seek, &fileLike_);
+}
+
+int AVIOFileLikeContext::read(void* opaque, uint8_t* buf, int buf_size) {
+  auto fileLike = static_cast<UniquePyObject*>(opaque);
+
+  // Note that we acquire the GIL outside of the loop. This is likely more
+  // efficient than releasing and acquiring it each loop iteration.
+  py::gil_scoped_acquire gil;
+
+  int totalNumRead = 0;
+  while (totalNumRead < buf_size) {
+    int request = buf_size - totalNumRead;
+
+    // The Python method returns the actual bytes, which we access through the
+    // py::bytes wrapper. That wrapper, however, does not provide us access to
+    // the underlying data pointer, which we need for the memcpy below. So we
+    // convert the bytes to a string_view to get access to the data pointer.
+    // Becauase it's a view and not a copy, it should be cheap.
+    auto bytesRead = static_cast<py::bytes>((*fileLike)->attr("read")(request));
+    auto bytesView = static_cast<std::string_view>(bytesRead);
+
+    int numBytesRead = static_cast<int>(bytesView.size());
+    if (numBytesRead == 0) {
+      break;
+    }
+
+    TORCH_CHECK(
+        numBytesRead <= request,
+        "Requested up to ",
+        request,
+        " bytes but, received ",
+        numBytesRead,
+        " bytes. The given object does not conform to read protocol of file object.");
+
+    std::memcpy(buf, bytesView.data(), numBytesRead);
+    buf += numBytesRead;
+    totalNumRead += numBytesRead;
+  }
+
+  return totalNumRead == 0 ? AVERROR_EOF : totalNumRead;
+}
+
+int64_t AVIOFileLikeContext::seek(void* opaque, int64_t offset, int whence) {
+  // We do not know the file size.
+  if (whence == AVSEEK_SIZE) {
+    return AVERROR(EIO);
+  }
+
+  auto fileLike = static_cast<UniquePyObject*>(opaque);
+  py::gil_scoped_acquire gil;
+  return py::cast<int64_t>((*fileLike)->attr("seek")(offset, whence));
+}
+
+} // namespace facebook::torchcodec

torchcodec/_core/AVIOFileLikeContext.h

@@ -0,0 +1,54 @@
+// 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.
+
+#pragma once
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include "src/torchcodec/_core/AVIOContextHolder.h"
+
+namespace py = pybind11;
+
+namespace facebook::torchcodec {
+
+// Enables uers to pass in a Python file-like object. We then forward all read
+// and seek calls back up to the methods on the Python object.
+class AVIOFileLikeContext : public AVIOContextHolder {
+ public:
+  explicit AVIOFileLikeContext(py::object fileLike);
+
+ private:
+  static int read(void* opaque, uint8_t* buf, int buf_size);
+  static int64_t seek(void* opaque, int64_t offset, int whence);
+
+  // Note that we dynamically allocate the Python object because we need to
+  // strictly control when its destructor is called. We must hold the GIL
+  // when its destructor gets called, as it needs to update the reference
+  // count. It's easiest to control that when it's dynamic memory. Otherwise,
+  // we'd have to ensure whatever enclosing scope holds the object has the GIL,
+  // and that's, at least, hard. For all of the common pitfalls, see:
+  //
+  //   https://pybind11.readthedocs.io/en/stable/advanced/misc.html#common-sources-of-global-interpreter-lock-errors
+  //
+  // We maintain a reference to the file-like object because the file-like
+  // object that was created on the Python side must live as long as our
+  // potential use. That is, even if there are no more references to the object
+  // on the Python side, we require that the object is still live.
+  struct PyObjectDeleter {
+    inline void operator()(py::object* obj) const {
+      if (obj) {
+        py::gil_scoped_acquire gil;
+        delete obj;
+      }
+    }
+  };
+
+  using UniquePyObject = std::unique_ptr<py::object, PyObjectDeleter>;
+  UniquePyObject fileLike_;
+};
+
+} // namespace facebook::torchcodec

torchcodec/_core/CMakeLists.txt

@@ -0,0 +1,237 @@
+cmake_minimum_required(VERSION 3.18)
+project(TorchCodec)
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+set(PYBIND11_FINDPYTHON ON)
+find_package(pybind11 REQUIRED)
+find_package(Torch REQUIRED)
+find_package(Python3 ${PYTHON_VERSION} EXACT COMPONENTS Development)
+
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra -pedantic -Werror ${TORCH_CXX_FLAGS}")
+
+function(make_torchcodec_sublibrary
+    library_name
+    type
+    sources
+    library_dependencies)
+
+    add_library(${library_name} ${type} ${sources})
+    set_target_properties(${library_name} PROPERTIES CXX_STANDARD 17)
+    target_include_directories(${library_name}
+        PRIVATE
+        ./../../../
+        "${TORCH_INSTALL_PREFIX}/include"
+        ${Python3_INCLUDE_DIRS}
+    )
+
+    # Avoid adding the "lib" prefix which we already add explicitly.
+    set_target_properties(${library_name} PROPERTIES PREFIX "")
+
+    target_link_libraries(
+        ${library_name}
+        PUBLIC
+        ${library_dependencies}
+    )
+endfunction()
+
+function(make_torchcodec_libraries
+    ffmpeg_major_version
+    ffmpeg_target)
+
+    # We create three shared libraries per version of FFmpeg, where the version
+    # is denoted by N:
+    #
+    # 1. libtorchcodec_decoderN.{ext}: Base library which contains the
+    #    implementation of VideoDecoder and everything VideoDecoder needs. On
+    #    Linux, {ext} is so. On Mac, it is dylib.
+    #
+    # 2. libtorchcodec_custom_opsN.{ext}: Implementation of the PyTorch custom
+    #    ops. Depends on libtorchcodec_decoderN.{ext}. On Linux, {ext} is so.
+    #    On Mac, it is dylib.
+    #
+    # 3. libtorchcodec_pybind_opsN.{ext}: Implementation of the pybind11 ops. We
+    #    keep these separate from the PyTorch custom ops because we have to
+    #    load these libraries separately on the Python side. Depends on
+    #    libtorchcodec_decoderN.{ext}. On BOTH Linux and Mac {ext} is so.
+
+    # 1. Create libtorchcodec_decoderN.{ext}.
+    set(decoder_library_name "libtorchcodec_decoder${ffmpeg_major_version}")
+    set(decoder_sources
+        AVIOContextHolder.cpp
+        FFMPEGCommon.cpp
+	DeviceInterface.cpp
+        SingleStreamDecoder.cpp
+        # TODO: lib name should probably not be "*_decoder*" now that it also
+        # contains an encoder
+        Encoder.cpp
+    )
+
+    if(ENABLE_CUDA)
+	    list(APPEND decoder_sources CudaDeviceInterface.cpp)
+    endif()
+
+    set(decoder_library_dependencies
+        ${ffmpeg_target}
+        ${TORCH_LIBRARIES}
+    )
+
+    if(ENABLE_CUDA)
+        list(APPEND decoder_library_dependencies
+            ${CUDA_nppi_LIBRARY}
+            ${CUDA_nppicc_LIBRARY}
+        )
+    endif()
+
+    make_torchcodec_sublibrary(
+        "${decoder_library_name}"
+        SHARED
+        "${decoder_sources}"
+        "${decoder_library_dependencies}"
+    )
+
+    # 2. Create libtorchcodec_custom_opsN.{ext}.
+    set(custom_ops_library_name "libtorchcodec_custom_ops${ffmpeg_major_version}")
+    set(custom_ops_sources
+        AVIOBytesContext.cpp
+        custom_ops.cpp
+    )
+    set(custom_ops_dependencies
+        ${decoder_library_name}
+        ${Python3_LIBRARIES}
+    )
+    make_torchcodec_sublibrary(
+        "${custom_ops_library_name}"
+        SHARED
+        "${custom_ops_sources}"
+        "${custom_ops_dependencies}"
+    )
+
+    # 3. Create libtorchcodec_pybind_opsN.so.
+    set(pybind_ops_library_name "libtorchcodec_pybind_ops${ffmpeg_major_version}")
+    set(pybind_ops_sources
+        AVIOFileLikeContext.cpp
+        pybind_ops.cpp
+    )
+    set(pybind_ops_dependencies
+       ${decoder_library_name}
+       pybind11::module # This library dependency makes sure we have the right
+                        # Python libraries included as well as all of the right
+                        # settings so that we can successfully load the shared
+                        # library as a Python module on Mac. If we instead use
+                        # ${Python3_LIBRARIES}, it works on Linux but not on
+                        # Mac.
+    )
+    make_torchcodec_sublibrary(
+        "${pybind_ops_library_name}"
+        MODULE # Note that this not SHARED; otherwise we build the wrong kind
+               # of library on Mac. On Mac, SHARED becomes .dylib and MODULE becomes
+               # a .so. We want pybind11 libraries to become .so. If this is
+               # changed to SHARED, we will be able to succesfully compile a
+               # .dylib, but we will not be able to succesfully import that as
+               # a Python module on Mac.
+        "${pybind_ops_sources}"
+        "${pybind_ops_dependencies}"
+    )
+    # pybind11 limits the visibility of symbols in the shared library to prevent
+    # stray initialization of py::objects. The rest of the object code must
+    # match. See:
+    #   https://pybind11.readthedocs.io/en/stable/faq.html#someclass-declared-with-greater-visibility-than-the-type-of-its-field-someclass-member-wattributes
+    target_compile_options(
+        ${pybind_ops_library_name}
+        PUBLIC
+      "-fvisibility=hidden"
+    )
+    # If we don't make sure this flag is set, we run into segfauls at import
+    # time on Mac. See:
+    #    https://github.com/pybind/pybind11/issues/3907#issuecomment-1170412764
+    target_link_options(
+        ${pybind_ops_library_name}
+        PUBLIC
+        "-undefined dynamic_lookup"
+    )
+
+    # Install all libraries.
+    set(
+        all_libraries
+        ${decoder_library_name}
+        ${custom_ops_library_name}
+        ${pybind_ops_library_name}
+    )
+
+    # The install step is invoked within CMakeBuild.build_library() in
+    # setup.py and just copies the built files from the temp
+    # cmake/setuptools build folder into the CMAKE_INSTALL_PREFIX folder. We
+    # still need to manually pass "DESTINATION ..." for cmake to copy those
+    # files in CMAKE_INSTALL_PREFIX instead of CMAKE_INSTALL_PREFIX/lib.
+    install(
+        TARGETS ${all_libraries}
+        LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}
+    )
+endfunction()
+
+if(DEFINED ENV{BUILD_AGAINST_ALL_FFMPEG_FROM_S3})
+    message(
+        STATUS
+        "Building and dynamically linking libtorchcodec against our pre-built
+        non-GPL FFmpeg libraries. These libraries are only used at build time,
+        you still need a different FFmpeg to be installed for run time!"
+    )
+
+    # This will expose the ffmpeg4, ffmpeg5, ffmpeg6, and ffmpeg7 targets
+    include(
+        ${CMAKE_CURRENT_SOURCE_DIR}/fetch_and_expose_non_gpl_ffmpeg_libs.cmake
+    )
+
+    make_torchcodec_libraries(7 ffmpeg7)
+    make_torchcodec_libraries(6 ffmpeg6)
+    make_torchcodec_libraries(4 ffmpeg4)
+    make_torchcodec_libraries(5 ffmpeg5)
+else()
+    message(
+        STATUS
+        "Building and dynamically linking libtorchcodec against the installed
+        FFmpeg libraries. This require pkg-config to be installed. If you have
+        installed FFmpeg from conda, make sure pkg-config is installed from
+        conda as well."
+    )
+    find_package(PkgConfig REQUIRED)
+    pkg_check_modules(LIBAV REQUIRED IMPORTED_TARGET
+        libavdevice
+        libavfilter
+        libavformat
+        libavcodec
+        libavutil
+        libswresample
+        libswscale
+    )
+
+    # Split libavcodec's version string by '.' and convert it to a list
+    string(REPLACE "." ";" libavcodec_version_list ${LIBAV_libavcodec_VERSION})
+    # Get the first element of the list, which is the major version
+    list(GET libavcodec_version_list 0 libavcodec_major_version)
+
+    if (${libavcodec_major_version} STREQUAL "58")
+        set(ffmpeg_major_version "4")
+    elseif (${libavcodec_major_version} STREQUAL "59")
+        set(ffmpeg_major_version "5")
+    elseif (${libavcodec_major_version} STREQUAL "60")
+        set(ffmpeg_major_version "6")
+    elseif (${libavcodec_major_version} STREQUAL "61")
+        set(ffmpeg_major_version "7")
+    else()
+        message(
+            FATAL_ERROR
+            "Unsupported libavcodec version: ${libavcodec_major_version}"
+        )
+    endif()
+
+    make_torchcodec_libraries(${ffmpeg_major_version} PkgConfig::LIBAV)
+
+    # Expose these values updwards so that the test compilation does not need
+    # to re-figure it out. FIXME: it's not great that we just copy-paste the
+    # library names.
+    set(libtorchcodec_library_name "libtorchcodec_decoder${ffmpeg_major_version}" PARENT_SCOPE)
+    set(libtorchcodec_custom_ops_name "libtorchcodec_custom_ops${ffmpeg_major_version}" PARENT_SCOPE)
+    set(libav_include_dirs ${LIBAV_INCLUDE_DIRS} PARENT_SCOPE)
+endif()