#pragma clang diagnostic push

#pragma clang diagnostic ignored "-Wexit-time-destructors"

#pragma clang diagnostic ignored "-Wdocumentation"

#pragma clang diagnostic ignored "-Wdeprecated-declarations"

#include <pybind11/pybind11.h>

#include <pybind11/stl.h>

#pragma clang diagnostic pop

#import <CoreML/CoreML.h>

#import <Availability.h>

#import <vector>

#import <optional>

#ifndef BUILT_WITH_MACOS15_SDK

#if TARGET_OS_OSX && defined(__MAC_15_0) && __MAC_OS_X_VERSION_MAX_ALLOWED >= __MAC_15_0

#define BUILT_WITH_MACOS15_SDK 1

#else

#define BUILT_WITH_MACOS15_SDK 0

#endif

#endif

#if BUILT_WITH_MACOS15_SDK

#pragma message ("Building with macOS 15+ SDK")

#else

#pragma message ("Building without macOS 15 SDK")

#endif

#if !defined(ML_COMPUTE_PLAN_IS_AVAILABLE) && __has_include(<CoreML/MLComputePlan.h>)

#define ML_COMPUTE_PLAN_IS_AVAILABLE 1

#endif

#if !defined(ML_MODEL_STRUCTURE_IS_AVAILABLE) && __has_include(<CoreML/MLModelStructure.h>)

#define ML_MODEL_STRUCTURE_IS_AVAILABLE 1

#endif

#if !defined(ML_COMPUTE_DEVICE_IS_AVAILABLE) && __has_include(<CoreML/MLComputeDeviceProtocol.h>)

#define ML_COMPUTE_DEVICE_IS_AVAILABLE 1

#endif

#if !defined(ML_MODEL_ASSET_IS_AVAILABLE) && __has_include(<CoreML/MLModelAsset.h>)

#define ML_MODEL_ASSET_IS_AVAILABLE 1

#endif

#if !defined(ML_STATE_IS_AVAILABLE) && __has_include(<CoreML/MLState.h>)

#define ML_STATE_IS_AVAILABLE 1

#endif

namespace py = pybind11;

namespace CoreML {

namespace Python {

struct State {

// MLState must be wrapped in a C++ class for PyBind.

inline State(id impl):

m_impl(impl) {}

#if ML_STATE_IS_AVAILABLE

API_AVAILABLE(macos(15.0))

inline MLState *getImpl() const {

return (MLState *)m_impl;

}

#endif

// Retrieves the value of the specified state variable from the model.

py::object readState(const std::string& stateName) const;

// Sets the value of the specified state variable in the model.

void writeState(const std::string& stateName, py::object value);

private:

// Type erase `m_impl` otherwise it will result in a compiler warning.

id m_impl = nil;

};

struct CPUComputeDevice {

inline CPUComputeDevice():

m_impl(nil) {}

// MLCPUComputeDevice must be wrapped in a C++ class for PyBind.

inline CPUComputeDevice(id impl):

m_impl(impl) {}

#if ML_COMPUTE_DEVICE_IS_AVAILABLE

API_AVAILABLE(macos(14.0))

inline MLCPUComputeDevice *getImpl() const {

return (MLCPUComputeDevice *)m_impl;

}

#endif

private:

// Type erase `m_impl` otherwise it will result in a compiler warning.

id m_impl = nil;

};

struct GPUComputeDevice {

inline GPUComputeDevice():

m_impl(nil) {}

// MLGPUComputeDevice must be wrapped in a C++ class for PyBind.

inline GPUComputeDevice(id impl):

m_impl(impl) {}

#if ML_COMPUTE_DEVICE_IS_AVAILABLE

API_AVAILABLE(macos(14.0))

inline MLGPUComputeDevice *getImpl() const {

return (MLGPUComputeDevice *)m_impl;

}

#endif

private:

// Type erase `m_impl` otherwise it will result in a compiler warning.

id m_impl = nil;

};

struct NeuralEngineComputeDevice {

inline NeuralEngineComputeDevice():

m_impl(nil) {}

// MLNeuralEngineComputeDevice must be wrapped in a C++ class for PyBind.

inline NeuralEngineComputeDevice(id impl):

m_impl(impl) {}

#if ML_COMPUTE_DEVICE_IS_AVAILABLE

API_AVAILABLE(macos(14.0))

inline MLNeuralEngineComputeDevice *getImpl() const {

return (MLNeuralEngineComputeDevice *)m_impl;

}

#endif

int getTotalCoreCount() const;

private:

// Type erase `m_impl` otherwise it will result in a compiler warning.

id m_impl = nil;

};

struct ModelStructureProgramOperation {

// MLModelStructureProgramOperation must be wrapped in a C++ class for PyBind.

inline ModelStructureProgramOperation(id impl):

m_impl(impl) {}

#if ML_MODEL_STRUCTURE_IS_AVAILABLE

API_AVAILABLE(macos(14.4))

inline MLModelStructureProgramOperation *getImpl() const {

return (MLModelStructureProgramOperation *)m_impl;

}

#endif

private:

// Type erase `m_impl` otherwise it will result in a compiler warning.

__weak id m_impl = nil;

};

struct ModelStructureNeuralNetworkLayer {

// ModelStructureNeuralNetworkLayer must be wrapped in a C++ class for PyBind.

inline ModelStructureNeuralNetworkLayer(id impl):

m_impl(impl) {}

#if ML_MODEL_STRUCTURE_IS_AVAILABLE

API_AVAILABLE(macos(14.4))

inline MLModelStructureNeuralNetworkLayer *getImpl() const {

return (MLModelStructureNeuralNetworkLayer *)m_impl;

}

#endif

private:

// Type erase `m_impl` otherwise it will result in a compiler warning.

__weak id m_impl = nil;

};

struct ComputePlan {

inline ComputePlan():

m_impl(nil), m_modelStructure(py::none()) {}

// MLComputePlan must be wrapped in a C++ class for PyBind.

inline ComputePlan(id impl, py::object modelStructure):

m_impl(impl),

m_modelStructure(modelStructure) {}

inline py::object getModelStructure() const {

return m_modelStructure;

}

#if ML_COMPUTE_PLAN_IS_AVAILABLE

API_AVAILABLE(macos(14.4))

inline MLComputePlan *getImpl() const {

return (MLComputePlan *)m_impl;

}

py::object getComputeDeviceUsageForMLProgramOperation(py::object operation);

py::object getComputeDeviceUsageForNeuralNetworkLayer(py::object layer);

py::object getEstimatedCostForMLProgramOperation(py::object operation);

#endif

private:

id m_impl = nil;

py::object m_modelStructure;

};

struct ModelAsset {

// MLModelAsset must be wrapped in a C++ class for PyBind.

inline ModelAsset(id impl, std::vector<py::bytes> datas):

m_impl(impl),

m_datas(std::move(datas)) {}

#if ML_MODEL_ASSET_IS_AVAILABLE

API_AVAILABLE(macos(13.0))

inline MLModelAsset *getImpl() const {

return (MLModelAsset *)m_impl;

}

#endif

id m_impl = nil;

std::vector<py::bytes> m_datas;

};

class Model {

private:

MLModel *m_model = nil;

NSURL *compiledUrl = nil;

bool m_deleteCompiledModelOnExit = false;

std::optional<uint64_t> m_loadDurationInNanoSeconds;

std::optional<uint64_t> m_lastPredictDurationInNanoSeconds;

public:

static py::bytes autoSetSpecificationVersion(const py::bytes& modelBytes);

static py::str compileModel(const std::string& urlStr);

static int32_t maximumSupportedSpecificationVersion();

static void setComputeUnit(MLModelConfiguration *configuration, const std::string& computeUnits);

Model(const Model&) = delete;

Model& operator=(const Model&) = delete;

~Model();

explicit Model(const std::string& urlStr,

const std::string& computeUnits,

const std::string& functionName,

const py::dict& optimizationHints,

const py::object& asset);

explicit Model(const std::string& urlStr,

const std::string& computeUnits,

const std::string& functionName,

const py::dict& optimizationHints);

explicit Model(MLModel* m_model, NSURL* compiledUrl, bool deleteCompiledModelOnExit);

py::list batchPredict(const py::list& batch);

py::str getCompiledModelPath() const;

py::dict predict(const py::dict& input, State* state=NULL);

py::object getLoadDurationInNanoSeconds() const;

py::object getLastPredictDurationInNanoSeconds() const;

#if BUILT_WITH_MACOS15_SDK

static void setOptimizationHints(MLModelConfiguration *configuration, const py::dict& optimizationHints);

State newState() const;

#endif

static py::object createModelAssetFromPath(const std::string& path);

static py::object createModelAssetFromMemory(const py::bytes& specData, const py::dict& blobMapping);

static py::object getModelStructure(const std::string& modelPath);

static py::list getAvailableComputeDevices();

static py::list getAllComputeDevices();

static py::object getComputePlan(const std::string& modelPath, const std::string& computeUnits);

};

}

}