Xinjie-Q's picture
Upload Mage-VL: unified codec-native streaming VLM (image+video understanding + proactive gate)
12acbba verified
Raw
History Blame Contribute Delete
1.66 kB
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
#include "def.h"
#include <torch/extension.h>
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
m.def("process_with_mask_cuda", &process_with_mask_cuda);
m.def("combine_for_reading_2x_cuda", &combine_for_reading_2x_cuda);
m.def("restore_y_2x_cuda", &restore_y_2x_cuda);
m.def("restore_y_4x_cuda", &restore_y_4x_cuda);
m.def("build_index_dec_cuda", &build_index_dec_cuda);
m.def("build_index_enc_cuda", &build_index_enc_cuda);
m.def("bias_quant_cuda", &bias_quant_cuda);
m.def("round_and_to_int8_cuda", &round_and_to_int8_cuda);
m.def("clamp_reciprocal_with_quant_cuda", &clamp_reciprocal_with_quant_cuda);
m.def("add_and_multiply_cuda", &add_and_multiply_cuda);
m.def("bias_pixel_shuffle_8_cuda", &bias_pixel_shuffle_8_cuda);
m.def("replicate_pad_cuda", &replicate_pad_cuda);
m.def("bias_wsilu_depthwise_conv2d_cuda", &bias_wsilu_depthwise_conv2d_cuda);
py::class_<DepthConvProxy>(m, "DepthConvProxy")
.def(py::init<>())
.def("set_param", &DepthConvProxy::set_param)
.def("set_param_with_adaptor", &DepthConvProxy::set_param_with_adaptor)
.def("forward", &DepthConvProxy::forward)
.def("forward_with_quant_step", &DepthConvProxy::forward_with_quant_step)
.def("forward_with_cat", &DepthConvProxy::forward_with_cat);
py::class_<SubpelConv2xProxy>(m, "SubpelConv2xProxy")
.def(py::init<>())
.def("set_param", &SubpelConv2xProxy::set_param)
.def("forward", &SubpelConv2xProxy::forward)
.def("forward_with_cat", &SubpelConv2xProxy::forward_with_cat);
}