model_runtime.cpp

#include "models/backend/model_runtime.h"

#include <algorithm>
#include <array>
#include <cstring>
#include <optional>
#include <limits>

#include <opencv2/imgproc.hpp>

#include "common/cv_utils.h"
#include "glog/logging.h"
#include "models/backend/request_geometry.h"

namespace mortred {
namespace models {
namespace backend {
namespace {

RuntimeStatus invalid_image(const std::string &operation, const std::string &message) {
    return runtime_error(StatusCode::MODEL_EMPTY_INPUT_IMAGE, "ImagePipeline " + operation + ": " + message);
}

bool set_error(StatusCode &status, std::string &destination, const RuntimeStatus &error) {
    if (status == StatusCode::OK) {
        status = error.status;
        destination = error.error;
    }
    return false;
}

bool valid_size(const cv::Size &size) { return size.width > 0 && size.height > 0; }

/*** float -> IEEE754 half, round-to-nearest-even */
inline std::uint16_t f32_to_f16_bits(float value) {
    std::uint32_t bits = 0;
    std::memcpy(&bits, &value, sizeof(bits));
    const std::uint16_t sign = static_cast<std::uint16_t>((bits >> 16) & 0x8000u);
    const std::uint32_t biased = (bits >> 23) & 0xFFu;
    const std::uint32_t mantissa = bits & 0x7FFFFFu;
    if (biased == 0xFFu) {
        return static_cast<std::uint16_t>(sign | 0x7C00u | (mantissa != 0 ? 0x200u : 0));
    }
    int32_t exponent = static_cast<int32_t>(biased) - 127 + 15;
    if (exponent >= 31) {
        return static_cast<std::uint16_t>(sign | 0x7C00u);
    }
    if (exponent <= 0) {
        if (exponent < -10) {
            return sign;
        }
        const std::uint32_t implicit = mantissa | 0x800000u;
        const int shift = 14 - exponent + 13;
        const std::uint32_t kept = implicit >> shift;
        const std::uint32_t rem = implicit & ((1u << shift) - 1u);
        const std::uint32_t rounded =
            kept + (((rem > (1u << (shift - 1))) || (rem == (1u << (shift - 1)) && (kept & 1u))) ? 1u : 0u);
        return static_cast<std::uint16_t>(sign | rounded);
    }
    std::uint32_t kept = mantissa >> 13;
    const std::uint32_t rem = mantissa & 0x1FFFu;
    kept += ((rem > 0x1000u) || (rem == 0x1000u && (kept & 1u))) ? 1u : 0u;
    if (kept == 0x400u) {
        kept = 0;
        ++exponent;
    }
    return static_cast<std::uint16_t>(sign | (static_cast<std::uint16_t>(exponent) << 10) | static_cast<std::uint16_t>(kept));
}

} // namespace

// Shallow, refcounted copy: no step writes into image_, so the only reason to
// deep-copy here was ownership, and cv::Mat already has that. Measured share
// of the preprocess chain is 13% (640x640), 21% (1080p), 30% (4K) — see
// scripts/microbench/clone_bench.cc.
ImagePipeline::ImagePipeline(const cv::Mat &image) : image_(image) {
    if (image.empty()) {
        set_error(status_, error_, invalid_image("constructor", "input image is empty"));
    }
}

bool ImagePipeline::require_private_float(const char *step) {
    // The constructor shares the caller's buffer. to_float() is the step that
    // replaces it with a private one, so the in-place steps below are only safe
    // after it; anything else would write through to the caller's image.
    // depth() not type(): the working image is CV_32FC3, whose type() encodes
    // the channel count as well.
    if (image_.depth() == CV_32F) {
        return true;
    }
    set_error(status_, error_,
              invalid_image(step, std::string("call to_float() before ") + step + "()"));
    return false;
}

ImagePipeline &ImagePipeline::bgr_to_rgb() {
    if (status_ != StatusCode::OK) {
        return *this;
    }
    if (image_.channels() != 3) {
        set_error(status_, error_, invalid_image("bgr_to_rgb", "expected a 3-channel image"));
        return *this;
    }
    cv::Mat converted;
    cv::cvtColor(image_, converted, cv::COLOR_BGR2RGB);
    image_ = std::move(converted);
    return *this;
}

ImagePipeline &ImagePipeline::bgra_to_rgb() {
    if (status_ != StatusCode::OK) {
        return *this;
    }
    if (image_.channels() != 4) {
        set_error(status_, error_, invalid_image("bgra_to_rgb", "expected a 4-channel image"));
        return *this;
    }
    cv::Mat converted;
    cv::cvtColor(image_, converted, cv::COLOR_BGRA2RGB);
    image_ = std::move(converted);
    return *this;
}

ImagePipeline &ImagePipeline::bgr_to_gray() {
    if (status_ != StatusCode::OK) {
        return *this;
    }
    if (image_.channels() != 3) {
        set_error(status_, error_, invalid_image("bgr_to_gray", "expected a 3-channel image"));
        return *this;
    }
    cv::Mat converted;
    cv::cvtColor(image_, converted, cv::COLOR_BGR2GRAY);
    image_ = std::move(converted);
    return *this;
}

ImagePipeline &ImagePipeline::rgb_to_bgr() {
    if (status_ != StatusCode::OK) {
        return *this;
    }
    if (image_.channels() != 3) {
        set_error(status_, error_, invalid_image("rgb_to_bgr", "expected a 3-channel image"));
        return *this;
    }
    cv::Mat converted;
    cv::cvtColor(image_, converted, cv::COLOR_RGB2BGR);
    image_ = std::move(converted);
    return *this;
}

ImagePipeline &ImagePipeline::resize(const cv::Size &size) {
    if (status_ != StatusCode::OK) {
        return *this;
    }
    if (!valid_size(size)) {
        set_error(status_, error_, invalid_image("resize", "target size must be positive"));
        return *this;
    }
    cv::Mat resized;
    cv::resize(image_, resized, size, 0.0, 0.0, cv::INTER_LINEAR);
    image_ = std::move(resized);
    return *this;
}

ImagePipeline &ImagePipeline::letterbox(const cv::Size &size, std::uint8_t pad_value) {
    if (status_ != StatusCode::OK) {
        return *this;
    }
    if (!valid_size(size) || !valid_size(image_.size())) {
        set_error(status_, error_, invalid_image("letterbox", "source and target size must be positive"));
        return *this;
    }
    const LetterboxGeometry geom = compute_letterbox_geometry(image_.size(), size);
    if (!valid_size(geom.unpadded)) {
        set_error(status_, error_, invalid_image("letterbox", "computed unpadded size is invalid"));
        return *this;
    }
    cv::Mat resized;
    if (image_.size() != geom.unpadded) {
        cv::resize(image_, resized, geom.unpadded, 0.0, 0.0, cv::INTER_LINEAR);
    } else {
        resized = image_;
    }
    const cv::Scalar pad = cv::Scalar::all(static_cast<double>(pad_value));
    cv::Mat padded;
    cv::copyMakeBorder(resized, padded, geom.pad_y, geom.pad_bottom, geom.pad_x, geom.pad_right, cv::BORDER_CONSTANT, pad);
    if (padded.size() != size) {
        set_error(status_, error_, invalid_image("letterbox", "padded size mismatches the network input"));
        return *this;
    }
    image_ = std::move(padded);
    return *this;
}

ImagePipeline &ImagePipeline::center_crop(const cv::Size &size) {
    if (status_ != StatusCode::OK) {
        return *this;
    }
    if (!valid_size(size) || size.width > image_.cols || size.height > image_.rows) {
        set_error(status_, error_, invalid_image("center_crop", "crop does not fit the source image"));
        return *this;
    }
    const int x = (image_.cols - size.width) / 2;
    const int y = (image_.rows - size.height) / 2;
    image_ = image_(cv::Rect(x, y, size.width, size.height)).clone();
    return *this;
}

ImagePipeline &ImagePipeline::to_float() {
    if (status_ != StatusCode::OK) {
        return *this;
    }
    cv::Mat converted;
    image_.convertTo(converted, CV_32F);
    image_ = std::move(converted);
    return *this;
}

ImagePipeline &ImagePipeline::scale(float factor) {
    if (status_ != StatusCode::OK) {
        return *this;
    }
    if (!require_private_float("scale")) {
        return *this;
    }
    image_ *= factor;
    return *this;
}

ImagePipeline &ImagePipeline::subtract(const std::array<float, 3> &values) {
    if (status_ != StatusCode::OK) {
        return *this;
    }
    if (image_.channels() != 3) {
        set_error(status_, error_, invalid_image("subtract", "expected a 3-channel image"));
        return *this;
    }
    if (!require_private_float("subtract")) {
        return *this;
    }
    cv::subtract(image_, cv::Scalar(values[0], values[1], values[2]), image_);
    return *this;
}

ImagePipeline &ImagePipeline::divide(const std::array<float, 3> &values) {
    if (status_ != StatusCode::OK) {
        return *this;
    }
    if (image_.channels() != 3) {
        set_error(status_, error_, invalid_image("divide", "expected a 3-channel image"));
        return *this;
    }
    if (!require_private_float("divide")) {
        return *this;
    }
    cv::divide(image_, cv::Scalar(values[0], values[1], values[2]), image_);
    return *this;
}

ImagePipeline &ImagePipeline::mean_std(const std::array<float, 3> &mean, const std::array<float, 3> &std) {
    subtract(mean);
    divide(std);
    return *this;
}

RuntimeResult<NamedTensor> ImagePipeline::pack(const std::string &name, bool nchw) const {
    if (status_ != StatusCode::OK) {
        return {status_, error_, {}};
    }
    if (name.empty()) {
        return {StatusCode::MODEL_EMPTY_INPUT_IMAGE, "ImagePipeline pack: tensor name is empty", {}};
    }
    if (image_.empty()) {
        return {StatusCode::MODEL_EMPTY_INPUT_IMAGE, "ImagePipeline pack: working image is empty", {}};
    }
    if (image_.type() != CV_32FC1 && image_.type() != CV_32FC3) {
        return {StatusCode::MODEL_EMPTY_INPUT_IMAGE, "ImagePipeline pack: expected CV_32FC1 or CV_32FC3", {}};
    }

    NamedTensor named;
    named.name = name;
    if (nchw) {
        named.tensor = Tensor::make<float>({1, image_.channels(), image_.rows, image_.cols});
        named.tensor.layout = TensorLayout::Nchw;
    } else {
        named.tensor = Tensor::make<float>({1, image_.rows, image_.cols, image_.channels()});
        named.tensor.layout = TensorLayout::Nhwc;
    }

    if (nchw) {
        const auto data = mortred::common::CvUtils::convert_to_chw_vec(image_);
        if (data.size() * sizeof(float) != named.tensor.byte_size()) {
            return {StatusCode::MODEL_EMPTY_INPUT_IMAGE, "ImagePipeline pack: CHW data size mismatch", {}};
        }
        std::memcpy(named.tensor.buffer.data(), data.data(), named.tensor.byte_size());
    } else {
        if (image_.isContinuous()) {
            const auto bytes = image_.total() * image_.elemSize();
            if (bytes != named.tensor.byte_size()) {
                return {StatusCode::MODEL_EMPTY_INPUT_IMAGE, "ImagePipeline pack: HWC data size mismatch", {}};
            }
            std::memcpy(named.tensor.buffer.data(), image_.data, bytes);
        } else {
            auto *destination = named.tensor.buffer.data();
            const size_t row_bytes = static_cast<size_t>(image_.cols) * image_.elemSize();
            for (int row = 0; row < image_.rows; ++row) {
                std::memcpy(destination, image_.ptr(row), row_bytes);
                destination += row_bytes;
            }
        }
    }
    return runtime_ok(std::move(named));
}

RuntimeResult<NamedTensor> ImagePipeline::nchw(const std::string &name) const { return nchw(name, DType::F32); }

RuntimeResult<NamedTensor> ImagePipeline::nchw(const std::string &name, DType dtype) const {
    auto packed = pack(name, true);
    if (!packed.ok()) {
        return packed;
    }
    return named_tensor_as(std::move(packed.value), dtype);
}

RuntimeResult<NamedTensor> ImagePipeline::nchw(const TensorInfo &input) const { return nchw(input.name, input.dtype); }

RuntimeResult<NamedTensor> ImagePipeline::nhwc(const std::string &name) const { return nhwc(name, DType::F32); }

RuntimeResult<NamedTensor> ImagePipeline::nhwc(const std::string &name, DType dtype) const {
    auto packed = pack(name, false);
    if (!packed.ok()) {
        return packed;
    }
    return named_tensor_as(std::move(packed.value), dtype);
}

RuntimeResult<NamedTensor> ImagePipeline::nhwc(const TensorInfo &input) const { return nhwc(input.name, input.dtype); }

RuntimeResult<NamedTensor> letterbox_bgr_f32_nchw(const cv::Mat &bgr, const cv::Size &network,
                                                  const std::string &tensor_name, std::uint8_t pad_value) {
    return letterbox_bgr_nchw(bgr, network, tensor_name, DType::F32, pad_value);
}

namespace {

/*** one pass over the source rows, writing the three plane rows at once
 * (the source is read once; plane 0 = R = BGR slot 2, 1 = G, 2 = B) */
template <typename T, typename CONVERT>
void fill_letterbox_planes(const cv::Mat &resized, const LetterboxGeometry &geom, int rows, int cols, T *base, T pad_value,
                           CONVERT convert) {
    const int band_y0 = geom.pad_y;
    const int band_y1 = rows - geom.pad_bottom;
    const int band_x0 = geom.pad_x;
    const int band_x1 = cols - geom.pad_right;
    const size_t plane_elems = static_cast<size_t>(rows) * static_cast<size_t>(cols);
    for (int y = 0; y < rows; ++y) {
        T *row_r = base + 0 * plane_elems + static_cast<size_t>(y) * cols;
        T *row_g = base + 1 * plane_elems + static_cast<size_t>(y) * cols;
        T *row_b = base + 2 * plane_elems + static_cast<size_t>(y) * cols;
        if (y < band_y0 || y >= band_y1) {
            std::fill(row_r, row_r + cols, pad_value);
            std::fill(row_g, row_g + cols, pad_value);
            std::fill(row_b, row_b + cols, pad_value);
            continue;
        }
        const auto *src = resized.ptr<const unsigned char>(y - band_y0);
        for (int x = 0; x < band_x0; ++x) {
            row_r[x] = pad_value;
            row_g[x] = pad_value;
            row_b[x] = pad_value;
        }
        for (int x = band_x0; x < band_x1; ++x) {
            const auto *px = src + static_cast<size_t>(x - band_x0) * 3;
            row_r[x] = convert(px[2]);
            row_g[x] = convert(px[1]);
            row_b[x] = convert(px[0]);
        }
        for (int x = band_x1; x < cols; ++x) {
            row_r[x] = pad_value;
            row_g[x] = pad_value;
            row_b[x] = pad_value;
        }
    }
}

}  // namespace

RuntimeResult<NamedTensor> letterbox_bgr_nchw(const cv::Mat &bgr, const cv::Size &network, const std::string &tensor_name,
                                              DType dtype, std::uint8_t pad_value) {
    if (bgr.empty() || bgr.type() != CV_8UC3) {
        return {StatusCode::MODEL_EMPTY_INPUT_IMAGE, "letterbox_bgr_nchw: expected a non-empty CV_8UC3 image", {}};
    }
    if (network.width <= 0 || network.height <= 0) {
        return {StatusCode::MODEL_EMPTY_INPUT_IMAGE, "letterbox_bgr_nchw: target size must be positive", {}};
    }
    if (tensor_name.empty()) {
        return {StatusCode::MODEL_EMPTY_INPUT_IMAGE, "letterbox_bgr_nchw: tensor name is empty", {}};
    }
    if (dtype != DType::F32 && dtype != DType::F16) {
        return {StatusCode::MODEL_EMPTY_INPUT_IMAGE, "letterbox_bgr_nchw: unsupported output dtype " +
                                                         std::string(dtype_to_string(dtype)), {}};
    }
    const LetterboxGeometry geom = compute_letterbox_geometry(bgr.size(), network);
    if (geom.unpadded.width <= 0 || geom.unpadded.height <= 0) {
        return {StatusCode::MODEL_EMPTY_INPUT_IMAGE, "letterbox_bgr_nchw: computed unpadded size is invalid", {}};
    }
    cv::Mat resized;
    if (bgr.size() != geom.unpadded) {
        cv::resize(bgr, resized, geom.unpadded, 0.0, 0.0, cv::INTER_LINEAR);
    } else {
        resized = bgr;
    }
    NamedTensor named;
    named.name = tensor_name;
    named.tensor = Tensor::make(dtype, {1, 3, static_cast<int64_t>(network.height), static_cast<int64_t>(network.width)});
    named.tensor.layout = TensorLayout::Nchw;
    const int rows = network.height;
    const int cols = network.width;
    const float scale_f = 1.0f / 255.0f;
    if (dtype == DType::F32) {
        // identical expression to the legacy chain: (float)px * (1/255f)
        const float pad_f = static_cast<float>(pad_value) * scale_f;
        fill_letterbox_planes<float>(
            resized, geom, rows, cols, reinterpret_cast<float *>(named.tensor.buffer.data()), pad_f,
            [scale_f](unsigned char px) { return static_cast<float>(px) * scale_f; });
    } else {
        // px*(1/255) has only 257 distinct values (256 pixel levels + pad),
        // so the fp16 conversion collapses into a one-shot LUT build and the
        // inner loop keeps doing a single L1 table read + 2-byte store
        std::array<std::uint16_t, 257> lut{};
        for (int level = 0; level < 256; ++level) {
            lut[static_cast<size_t>(level)] = f32_to_f16_bits(static_cast<float>(level) * scale_f);
        }
        lut[256] = f32_to_f16_bits(static_cast<float>(pad_value) * scale_f);
        fill_letterbox_planes<std::uint16_t>(
            resized, geom, rows, cols, reinterpret_cast<std::uint16_t *>(named.tensor.buffer.data()), lut[256],
            [&lut](unsigned char px) { return lut[px]; });
    }
    return runtime_ok(std::move(named));
}

RuntimeResult<NamedTensor> named_tensor_as(NamedTensor tensor, DType dtype) {
    if (tensor.tensor.dtype == dtype) {
        return runtime_ok(std::move(tensor));
    }
    if (tensor.tensor.dtype != DType::F32 || dtype != DType::F16) {
        return {StatusCode::MODEL_EMPTY_INPUT_IMAGE,
                std::string("named_tensor_as: cannot cast ") + dtype_to_string(tensor.tensor.dtype) + " to " +
                    dtype_to_string(dtype),
                {}};
    }
    Tensor out = Tensor::make(DType::F16, tensor.tensor.shape);
    out.layout = tensor.tensor.layout;
    const int64_t count = tensor.tensor.element_count();
    if (count < 0) {
        return {StatusCode::MODEL_EMPTY_INPUT_IMAGE, "named_tensor_as: invalid element count", {}};
    }
    const auto *src = reinterpret_cast<const float *>(tensor.tensor.buffer.data());
    auto *dst = reinterpret_cast<std::uint16_t *>(out.buffer.data());
    for (int64_t i = 0; i < count; ++i) {
        dst[static_cast<size_t>(i)] = f32_to_f16_bits(src[static_cast<size_t>(i)]);
    }
    tensor.tensor = std::move(out);
    return runtime_ok(std::move(tensor));
}

RuntimeResult<cv::Mat> ImagePipeline::mat() const {
    if (status_ != StatusCode::OK) {
        return {status_, error_, {}};
    }
    return runtime_ok(image_.clone());
}

OutputReader::OutputReader(const std::vector<NamedTensor> &outputs, std::string name) : outputs_(&outputs), name_(std::move(name)) {}

OutputReader &OutputReader::f32() {
    contract_.dtype = DType::F32;
    return *this;
}

OutputReader &OutputReader::shape(std::vector<int64_t> shape) {
    contract_.rank = shape.size();
    contract_.shape = std::move(shape);
    return *this;
}

OutputReader &OutputReader::nchw() {
    contract_.layout = TensorLayout::Nchw;
    return *this;
}

OutputReader &OutputReader::finite() {
    require_finite_ = true;
    return *this;
}

RuntimeResult<F32OutputView> OutputReader::read() const {
    F32OutputView view;
    const auto status =
        validated_f32_named_output(*outputs_, name_, contract_, "OutputReader", &view, require_finite_);
    if (status != StatusCode::OK) {
        return {status, "OutputReader contract failed", {}};
    }
    return runtime_ok(std::move(view));
}

ParamReader::ParamReader(const toml::table &params, std::string log_prefix) : params_(params), prefix_(std::move(log_prefix)) {}

ParamReader &ParamReader::fail(const std::string &message) {
    if (status_.ok()) {
        status_ = runtime_error(StatusCode::MODEL_INIT_FAILED, prefix_ + ": " + message);
        LOG(ERROR) << status_.error;
    }
    return *this;
}

ParamReader &ParamReader::get(const std::string &key, int32_t *value) {
    last_key_ = key;
    if (status_.status != StatusCode::OK) {
        return *this;
    }
    if (!params_.contains(key) || !params_.at(key).is_integer()) {
        return fail("param '" + key + "' must be an integer");
    }
    const auto parsed = params_.at(key).value<int64_t>();
    if (!parsed.has_value()) {
        return fail("param '" + key + "' must be an integer");
    }
    if (*parsed < std::numeric_limits<int32_t>::min() || *parsed > std::numeric_limits<int32_t>::max()) {
        return fail("param '" + key + "' must fit in int32");
    }
    *value = static_cast<int32_t>(*parsed);
    return *this;
}

ParamReader &ParamReader::get(const std::string &key, int64_t *value) {
    last_key_ = key;
    if (!status_.ok()) {
        return *this;
    }
    if (!params_.contains(key) || !params_.at(key).is_integer()) {
        return fail("param '" + key + "' must be an integer");
    }
    const auto parsed = params_.at(key).value<int64_t>();
    if (!parsed.has_value()) {
        return fail("param '" + key + "' must be an integer");
    }
    *value = *parsed;
    return *this;
}

ParamReader &ParamReader::get(const std::string &key, float *value) {
    double parsed = 0.0;
    const auto status = get(key, &parsed);
    if (status.ok()) {
        *value = static_cast<float>(parsed);
    }
    return *this;
}

ParamReader &ParamReader::get(const std::string &key, double *value) {
    last_key_ = key;
    if (!status_.ok()) {
        return *this;
    }
    if (!params_.contains(key) || (!params_.at(key).is_floating_point() && !params_.at(key).is_integer())) {
        return fail("param '" + key + "' must be numeric");
    }
    const auto parsed = params_.at(key).value<double>();
    if (!parsed.has_value()) {
        return fail("param '" + key + "' must be numeric");
    }
    *value = *parsed;
    return *this;
}

ParamReader &ParamReader::get(const std::string &key, bool *value) {
    last_key_ = key;
    if (!status_.ok()) {
        return *this;
    }
    if (!params_.contains(key) || !params_.at(key).is_boolean()) {
        return fail("param '" + key + "' must be boolean");
    }
    const auto parsed = params_.at(key).value<bool>();
    if (!parsed.has_value()) {
        return fail("param '" + key + "' must be boolean");
    }
    *value = *parsed;
    return *this;
}

ParamReader &ParamReader::get(const std::string &key, std::string *value) {
    last_key_ = key;
    if (!status_.ok()) {
        return *this;
    }
    if (!params_.contains(key) || !params_.at(key).is_string()) {
        return fail("param '" + key + "' must be a string");
    }
    const auto parsed = params_.at(key).value<std::string>();
    if (!parsed.has_value()) {
        return fail("param '" + key + "' must be a string");
    }
    *value = *parsed;
    return *this;
}

ParamReader &ParamReader::get(const std::string &key, cv::Size *value) {
    last_key_ = key;
    if (!status_.ok()) {
        return *this;
    }
    const toml::array *array = params_.contains(key) ? params_.at(key).as_array() : nullptr;
    if (array == nullptr || array->size() != 2) {
        return fail("param '" + key + "' must be [height, width]");
    }
    const auto height = (*array)[0].value<int64_t>();
    const auto width = (*array)[1].value<int64_t>();
    if (!height.has_value() || !width.has_value() || *height <= 0 || *width <= 0) {
        return fail("param '" + key + "' must contain positive integer dimensions");
    }
    if (*height > std::numeric_limits<int>::max() || *width > std::numeric_limits<int>::max()) {
        return fail("param '" + key + "' dimensions must fit in int");
    }
    *value = cv::Size(static_cast<int>(*width), static_cast<int>(*height));
    return *this;
}

ParamReader &ParamReader::min(double value) {
    if (!status_.ok()) {
        return *this;
    }
    const auto parsed = params_.contains(last_key_) ? params_.at(last_key_).value<double>() : std::optional<double>{};
    if (parsed.has_value() && *parsed < value) {
        fail("param '" + last_key_ + "' must be >= " + std::to_string(value));
    }
    return *this;
}

ParamReader &ParamReader::max(double value) {
    if (!status_.ok()) {
        return *this;
    }
    const auto parsed = params_.contains(last_key_) ? params_.at(last_key_).value<double>() : std::optional<double>{};
    if (parsed.has_value() && *parsed > value) {
        fail("param '" + last_key_ + "' must be <= " + std::to_string(value));
    }
    return *this;
}

ParamReader &ParamReader::non_empty() {
    if (!status_.ok()) {
        return *this;
    }
    const auto parsed = params_.contains(last_key_) ? params_.at(last_key_).value<std::string>() : std::optional<std::string>{};
    if (parsed.has_value() && parsed->empty()) {
        fail("param '" + last_key_ + "' must not be empty");
    }
    return *this;
}

ParamReader &ParamReader::array_size(size_t size) {
    if (!status_.ok()) {
        return *this;
    }
    const auto *array = params_.contains(last_key_) ? params_.at(last_key_).as_array() : nullptr;
    if (array != nullptr && array->size() != size) {
        fail("param '" + last_key_ + "' must have " + std::to_string(size) + " elements");
    }
    return *this;
}

ParamReader &ParamReader::allow_only_keys(const std::vector<std::string> &keys) {
    if (!status_.ok()) {
        return *this;
    }
    for (const auto &item : params_) {
        const std::string key(item.first.str());
        if (std::find(keys.begin(), keys.end(), key) == keys.end()) {
            fail("unknown param '" + key + "'");
        }
    }
    return *this;
}

SessionIoValidator::SessionIoValidator(const InferenceSession &session) : session_(session) {}

SessionIoValidator &SessionIoValidator::input(const std::string &name) {
    use_output_ = false;
    name_ = name;
    return *this;
}

SessionIoValidator &SessionIoValidator::output(const std::string &name) {
    use_output_ = true;
    name_ = name;
    return *this;
}

SessionIoValidator &SessionIoValidator::dtype(DType value) {
    dtype_ = value;
    return *this;
}

SessionIoValidator &SessionIoValidator::f32() { return dtype(DType::F32); }

SessionIoValidator &SessionIoValidator::allow_fp16() {
    allow_fp16_ = true;
    return *this;
}

SessionIoValidator &SessionIoValidator::rank(size_t value) {
    rank_ = value;
    return *this;
}

SessionIoValidator &SessionIoValidator::shape(std::vector<int64_t> value) {
    expected_shape_ = std::move(value);
    rank_ = expected_shape_.size();
    return *this;
}

SessionIoValidator &SessionIoValidator::nchw() {
    has_layout_ = true;
    nchw_ = true;
    return *this;
}

SessionIoValidator &SessionIoValidator::nhwc() {
    has_layout_ = true;
    nchw_ = false;
    return *this;
}

SessionIoValidator &SessionIoValidator::channels(int64_t value) {
    channels_ = value;
    return *this;
}

SessionIoValidator &SessionIoValidator::static_shape() {
    static_shape_ = true;
    return *this;
}

SessionIoValidator &SessionIoValidator::allow_dynamic_batch() {
    allow_dynamic_batch_ = true;
    return *this;
}

RuntimeResult<TensorInfo> SessionIoValidator::validate() const {
    const auto &infos = use_output_ ? session_.outputs() : session_.inputs();
    const auto found = name_.empty()
                           ? infos.begin()
                           : std::find_if(infos.begin(), infos.end(), [this](const TensorInfo &item) { return item.name == name_; });
    if (found == infos.end()) {
        const std::string kind = use_output_ ? "output" : "input";
        const std::string name = name_.empty() ? "<first>" : name_;
        return {StatusCode::MODEL_INIT_FAILED, "session " + kind + " '" + name + "' is missing", {}};
    }
    const auto &info = *found;
    const std::string description = info.to_string();
    if (info.dtype != dtype_) {
        const bool fp16_ok = allow_fp16_ && dtype_ == DType::F32 && info.dtype == DType::F16;
        if (!fp16_ok) {
            return {StatusCode::MODEL_INIT_FAILED, "unexpected session io dtype: " + description, {}};
        }
    }
    if (rank_ != 0 && info.shape.size() != rank_) {
        return {StatusCode::MODEL_INIT_FAILED, "unexpected session io rank: " + description, {}};
    }
    if (expected_shape_.size() != 0) {
        if (expected_shape_.size() != info.shape.size()) {
            return {StatusCode::MODEL_INIT_FAILED, "unexpected session io rank: " + description, {}};
        }
        for (size_t idx = 0; idx < expected_shape_.size(); ++idx) {
            if (expected_shape_[idx] >= 0 && info.shape[idx] != expected_shape_[idx]) {
                return {StatusCode::MODEL_INIT_FAILED, "unexpected session io shape: " + description, {}};
            }
        }
    }
    if (channels_ >= 0) {
        if (info.shape.size() != 4 || (nchw_ ? info.shape[1] : info.shape[3]) != channels_) {
            return {StatusCode::MODEL_INIT_FAILED, "unexpected session io channels: " + description, {}};
        }
    }
    if (has_layout_ && info.layout != TensorLayout::Unknown) {
        const auto expected = nchw_ ? TensorLayout::Nchw : TensorLayout::Nhwc;
        if (info.layout != expected) {
            return {StatusCode::MODEL_INIT_FAILED, "unexpected session io layout: " + description, {}};
        }
    }
    if (static_shape_ && info.dynamic) {
        return {StatusCode::MODEL_INIT_FAILED, "dynamic session io is not allowed: " + description, {}};
    }
    if (allow_dynamic_batch_) {
        for (size_t idx = 1; idx < info.shape.size(); ++idx) {
            if (info.shape[idx] <= 0) {
                return {StatusCode::MODEL_INIT_FAILED, "dynamic non-batch dimension: " + description, {}};
            }
        }
    }
    return runtime_ok(info);
}

} // namespace backend
} // namespace models
} // namespace mortred