Error

Author: codeSamuraii

CMakeLists.txt

@@ -21,6 +21,30 @@ endif()
 list(APPEND CMAKE_PREFIX_PATH "${NB_DIR}")
 find_package(nanobind CONFIG REQUIRED)
 
+# Find MLX
+execute_process(
+  COMMAND "${Python_EXECUTABLE}" -c "import mlx; print(mlx.__path__[0])"
+  OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE MLX_PYTHON_DIR
+  RESULT_VARIABLE MLX_RESULT)
+
+if(MLX_RESULT)
+  message(FATAL_ERROR "MLX not found. Install with: pip install mlx")
+endif()
+
+# MLX include and library paths
+set(MLX_INCLUDE_DIR "${MLX_PYTHON_DIR}/include")
+set(MLX_LIB_DIR "${MLX_PYTHON_DIR}/lib")
+
+# Find MLX library
+find_library(MLX_LIBRARY
+  NAMES mlx
+  PATHS "${MLX_LIB_DIR}"
+  NO_DEFAULT_PATH
+  REQUIRED)
+
+message(STATUS "MLX include dir: ${MLX_INCLUDE_DIR}")
+message(STATUS "MLX library: ${MLX_LIBRARY}")
+
 # Apple frameworks
 find_library(AVFOUNDATION_FRAMEWORK AVFoundation REQUIRED)
 find_library(COREVIDEO_FRAMEWORK CoreVideo REQUIRED)
@@ -42,9 +66,13 @@ nanobind_add_module(
   ${SOURCES}
 )
 
-target_include_directories(_viteo PRIVATE src/native/include)
+target_include_directories(_viteo PRIVATE
+  src/native/include
+  ${MLX_INCLUDE_DIR}
+)
 
 target_link_libraries(_viteo PRIVATE
+  ${MLX_LIBRARY}
   ${AVFOUNDATION_FRAMEWORK}
   ${COREVIDEO_FRAMEWORK}
   ${COREMEDIA_FRAMEWORK}
@@ -54,6 +82,15 @@ target_link_libraries(_viteo PRIVATE
   ${COREIMAGE_FRAMEWORK}
 )
 
+# Set rpath to find MLX library at runtime
+# Use @loader_path to find MLX relative to the module location
+set_target_properties(_viteo PROPERTIES
+  BUILD_RPATH "${MLX_LIB_DIR}"
+  INSTALL_RPATH "@loader_path/mlx/lib"
+  BUILD_WITH_INSTALL_RPATH FALSE
+  INSTALL_RPATH_USE_LINK_PATH TRUE
+)
+
 # Enable Objective-C++ for .mm files
 set_source_files_properties(
   src/native/src/frame_extractor.mm

src/native/src/frame_extractor.h src/native/include/frame_extractor.hsrc/native/src/frame_extractor.h renamed to src/native/include/frame_extractor.h

@@ -4,6 +4,7 @@
 #include <string>
 #include <cstdint>
 #include <vector>
+#include <mlx/mlx.h>
 
 namespace viteo {
 
@@ -16,7 +17,7 @@ class FrameExtractor {
     /// Open video file for extraction
     bool open(const std::string& path);
 
-    /// Get next frame as BGRA data (returns nullptr when done)
+    /// Get next frame as MLX array (returns empty array when done)
     uint8_t* next_frame();
 
     /// Reset to beginning or specific frame index

src/native/src/bindings.cpp

@@ -7,23 +7,17 @@ namespace nb = nanobind;
 using namespace viteo;
 
 /// Create MLX array from raw BGRA buffer
-nb::object create_mlx_array(uint8_t* data, int height, int width) {
-    if (!data) return nb::none();
+mlx::core::array create_mlx_array(uint8_t* data, int height, int width) {
+    auto arr = mlx::core::array(
+        data,
+        mlx::core::Shape{ (int32_t)height, (int32_t)width, int32_t(4) },
+        mlx::core::uint8
+    );
 
-    // Import MLX
-    nb::object mlx = nb::module_::import_("mlx.core");
-    nb::object mx_array = mlx.attr("array");
-    nb::object mx_uint8 = mlx.attr("uint8");
+    // Eval the array
+    mlx::core::eval({arr});
 
-    // Create memory view
-    size_t size = height * width * 4;
-    nb::object memview = nb::steal(PyMemoryView_FromMemory(
-        reinterpret_cast<char*>(data), size, PyBUF_READ
-    ));
-
-    // Create MLX array and reshape
-    nb::object arr = mx_array(memview, mx_uint8);
-    return arr.attr("reshape")(nb::make_tuple(height, width, 4));
+    return arr;
 }
 
 NB_MODULE(_viteo, m) {
@@ -34,13 +28,12 @@ NB_MODULE(_viteo, m) {
         .def("open", &FrameExtractor::open, nb::arg("path"),
             "Open video file for extraction")
         .def("next_frame",
-            [](FrameExtractor& self) -> nb::object {
+            [](FrameExtractor& self) -> mlx::core::array {
                 uint8_t* frame_data;
                 {
                     nb::gil_scoped_release release;
                     frame_data = self.next_frame();
                 }
-                if (!frame_data) return nb::none();
                 return create_mlx_array(frame_data, self.height(), self.width());
             },
             "Get next frame as MLX array (None when done)")
@@ -52,7 +45,7 @@ NB_MODULE(_viteo, m) {
         .def_prop_ro("total_frames", &FrameExtractor::total_frames, "Total frames")
         .def("__iter__", [](nb::object self) { return self; })
         .def("__next__",
-            [](FrameExtractor& self) -> nb::object {
+            [](FrameExtractor& self) -> mlx::core::array {
                 uint8_t* frame_data;
                 {
                     nb::gil_scoped_release release;

src/native/src/frame_extractor.mm

@@ -132,6 +132,7 @@ bool open(const std::string& path) {
             isOpen = true;
             if (!setupReader(0)) return false;
 
+            DEBUG_LOG("Video opened successfully");
             return prefetchFrame();
         }
     }
@@ -233,22 +234,26 @@ void copyFrameData(CVImageBufferRef imageBuffer, uint8_t* dst) {
 
     bool prefetchFrame() {
         if (!isOpen || !reader || !output) {
+            DEBUG_LOG("Cannot prefetch frame, extractor not open");
             return false;
         }
 
         @autoreleasepool {
             if (reader.status != AVAssetReaderStatusReading) {
+                DEBUG_LOG("Reader not in reading state");
                 return false;
             }
 
             CMSampleBufferRef sampleBuffer = [output copyNextSampleBuffer];
             if (!sampleBuffer) {
+                DEBUG_LOG("No more samples available");
                 return false;
             }
 
             CVImageBufferRef imageBuffer = CMSampleBufferGetImageBuffer(sampleBuffer);
             if (!imageBuffer) {
                 CFRelease(sampleBuffer);
+                DEBUG_LOG("Failed to get image buffer from sample");
                 return false;
             }
 
@@ -258,6 +263,7 @@ bool prefetchFrame() {
 
             CFRelease(sampleBuffer);
             has_prefetched_frame = true;
+            DEBUG_LOG("Prefetched frame " << currentFrame);
 
             return true;
         }
@@ -269,11 +275,13 @@ bool prefetchFrame() {
             return nullptr;
         }
 
+        DEBUG_LOG("Swapping frame buffers for frame " << currentFrame);
         std::swap(frame_buffer, prefetch_buffer);
         currentFrame++;
 
         prefetchFrame();
 
+        DEBUG_LOG("Returning frame buffer " << (currentFrame - 1));
         return frame_buffer.data();
     }
 

src/viteo/__init__.py

@@ -16,6 +16,7 @@
         process_frame(frame)
 """
 import pathlib
+import mlx.core as mx
 from _viteo import FrameExtractor as _FrameExtractor
 from typing import Optional