diff --git a/express/Expr.cpp b/express/Expr.cpp index 4080b7a9b6..0e03c26da3 100644 --- a/express/Expr.cpp +++ b/express/Expr.cpp @@ -1049,13 +1049,25 @@ std::vector Variable::load(const uint8_t* buffer, size_t length) { // Set tensor shape from net expr->mCanDecompose = false; } - if (nullptr != expr->get() || expr->inputType() == VARP::INPUT) { + if (nullptr != expr->get()) { + // Non-INPUT node: keep the existing behavior (replace with cloned tensor from net metadata) for (int index = 0; index < op->outputIndexes.size(); ++index) { auto outputIndex = op->outputIndexes[index]; delete expr->inside()->mOutputTensors[index]; expr->inside()->mOutputTensors[index] = Tensor::clone(allTensors[outputIndex].get()); Utils::copyTensorToInfo(expr->inside()->mOutputInfos.data() + index, expr->inside()->mOutputTensors[index]); } + } else if (expr->inputType() == VARP::INPUT) { + // INPUT node: keep the tensor created by Expr::create (host already allocated, + // usage = INPUT, memoryType = MEMORY_HOST). Only migrate quantAttr if the model + // declares quantization info for this input. + for (int index = 0; index < op->outputIndexes.size(); ++index) { + auto outputIndex = op->outputIndexes[index]; + auto srcTensor = allTensors[outputIndex].get(); + if (nullptr != srcTensor && TensorUtils::getDescribe(srcTensor)->quantAttr) { + TensorUtils::getDescribe(expr->inside()->mOutputTensors[index])->quantAttr = TensorUtils::getDescribe(srcTensor)->quantAttr; + } + } } for (int index = 0; index < op->outputIndexes.size(); ++index) { diff --git a/source/backend/arm82/Arm82Functions.cpp b/source/backend/arm82/Arm82Functions.cpp index 2125db807c..9381027bc0 100644 --- a/source/backend/arm82/Arm82Functions.cpp +++ b/source/backend/arm82/Arm82Functions.cpp @@ -2988,7 +2988,7 @@ bool Arm82Functions::init() { gInstance->supportSDot = origin->supportSDot; gInstance->supportI8mm = origin->supportI8mm; gInstance->supportSME2 = origin->supportSME2; -#if defined(MNN_SME2) && defined(__aarch64__) && defined(MNN_USE_NEON) +#if defined(MNN_SME2) && defined(MNN_SUPPORT_TRANSFORMER_FUSE) && defined(__aarch64__) && defined(MNN_USE_NEON) gInstance->supportFp16FML = origin->supportFp16FML; #endif gInstance->smeCoreNumber = origin->smeCoreNumber; diff --git a/source/backend/cpu/compute/Convolution1x1Strassen.cpp b/source/backend/cpu/compute/Convolution1x1Strassen.cpp index 2460d5df1d..b35342026d 100644 --- a/source/backend/cpu/compute/Convolution1x1Strassen.cpp +++ b/source/backend/cpu/compute/Convolution1x1Strassen.cpp @@ -119,8 +119,8 @@ ErrorCode Convolution1x1Strassen::onResize(const std::vector &inputs, } unit.offset[1] = 0; unit.offset[2] = 0; - unit.offset[0] = core->pack * planeStart * bytes; - unit.offset[3] = core->pack * planeStart * bytes; + unit.offset[0] = static_cast(core->pack) * planeStart * bytes; + unit.offset[3] = static_cast(core->pack) * planeStart * bytes; unit.mStracssenComputor.reset(new StrassenMatrixComputor(backend(), maxDepth)); int e = planeSize; int l = ic; @@ -162,19 +162,20 @@ ErrorCode Convolution1x1Strassen::onResize(const std::vector &inputs, } auto ocStartWeight = (ocStart * core->pack) / hPack; auto ocWeightSize = std::min(UP_DIV((ocSize * core->pack), hPack), mResource->mWeight->length(0) - ocStartWeight); - unit.offset[1] = hPack * icAlign * ocStartWeight * mWeightBytes; - unit.offset[2] = core->pack * ocStart * bytes; + // Large output projections can exceed 2 GB. Promote before multiplying to avoid int overflow. + unit.offset[1] = static_cast(hPack) * icAlign * ocStartWeight * mWeightBytes; + unit.offset[2] = static_cast(core->pack) * ocStart * bytes; unit.offset[0] = 0; - unit.offset[3] = core->pack * matrixSizeE * ocStart * bytes; + unit.offset[3] = static_cast(core->pack) * matrixSizeE * ocStart * bytes; unit.mStracssenComputor.reset(new StrassenMatrixComputor(backend(), maxDepth)); int e = matrixSizeE; int l = ic; int h = std::min(ocSize * core->pack, ocWeightSize * hPack); uint8_t* aPtr = nullptr; - auto bPtr = TensorUtils::getDescribeOrigin(mResource->mWeight.get())->mem->chunk() + hPack * icAlign * ocStartWeight * mWeightBytes; + auto bPtr = TensorUtils::getDescribeOrigin(mResource->mWeight.get())->mem->chunk() + unit.offset[1]; uint8_t* cPtr = nullptr; - auto biasPtr = TensorUtils::getDescribeOrigin(mResource->mBias.get())->mem->chunk() + core->pack * ocStart * bytes; + auto biasPtr = TensorUtils::getDescribeOrigin(mResource->mBias.get())->mem->chunk() + unit.offset[2]; memoryPool->beginGroup(); auto code = unit.mStracssenComputor->onEncode(e, l, h, matrixSizeE * core->pack, UP_DIV(l, lPack) * lPack * hPack, matrixSizeE * core->pack, aPtr, bPtr, cPtr, true, biasPtr, postParameters); if (NO_ERROR != code) { diff --git a/source/backend/cpu/compute/Convolution1x1Strassen.hpp b/source/backend/cpu/compute/Convolution1x1Strassen.hpp index 531523e720..8d5bbfa3e5 100644 --- a/source/backend/cpu/compute/Convolution1x1Strassen.hpp +++ b/source/backend/cpu/compute/Convolution1x1Strassen.hpp @@ -29,7 +29,7 @@ class Convolution1x1Strassen : public CPUConvolution { struct Unit { bool mValid = true; - int offset[4];//Input, Weight, Output, Bias + size_t offset[4]; // Input, Weight, Output, Bias std::shared_ptr mStracssenComputor; }; diff --git a/source/backend/cpu/compute/ImageProcessFunction.cpp b/source/backend/cpu/compute/ImageProcessFunction.cpp index 6a2c057d08..ae2e69d60d 100644 --- a/source/backend/cpu/compute/ImageProcessFunction.cpp +++ b/source/backend/cpu/compute/ImageProcessFunction.cpp @@ -273,7 +273,7 @@ void MNNC3ToXYZ(const unsigned char* source, unsigned char* dest, size_t count, C3 = coeffs[r1], C4 = coeffs[4], C5 = coeffs[b1], C6 = coeffs[r2], C7 = coeffs[7], C8 = coeffs[b2]; int sta = 0; - + #if defined MNN_USE_NEON int countD8 = (int)count / 8; if (countD8 > 0) { @@ -282,7 +282,7 @@ void MNNC3ToXYZ(const unsigned char* source, unsigned char* dest, size_t count, sta = countD8 * 8; } #endif - + for (int i = sta; i < count; ++i) { int r = source[3 * i + 0]; int g = source[3 * i + 1]; @@ -333,7 +333,7 @@ void MNNC3ToBGR555(const unsigned char* source, unsigned char* dest, size_t coun } else { MNNRGBToBGR555Fast(source, dest, countD8); } - + i = countD8 * 8; } #endif @@ -755,8 +755,8 @@ static void _sampleBilinearCommon(const unsigned char* source, unsigned char* de float x = __clamp(curPoints.fX, 0, xMax); int y0 = (int)y; int x0 = (int)x; - int y1 = (int)ceilf(y); - int x1 = (int)ceilf(x); + int y1 = std::min((int)ceilf(y), (int)(ih - 1)); + int x1 = std::min((int)ceilf(x), (int)(iw - 1)); float xF = x - (float)x0; float yF = y - (float)y0; diff --git a/source/backend/cpu/riscv/rvv/MNNAccumulateSequenceNumber.cpp b/source/backend/cpu/riscv/rvv/MNNAccumulateSequenceNumber.cpp index 9e13ba33f2..851b07b080 100644 --- a/source/backend/cpu/riscv/rvv/MNNAccumulateSequenceNumber.cpp +++ b/source/backend/cpu/riscv/rvv/MNNAccumulateSequenceNumber.cpp @@ -1,17 +1,16 @@ #include + void MNNAccumulateSequenceNumber_RVV(float* dst, const float* src, int size) { - size_t vl = __riscv_vsetvlmax_e32m1(); - vfloat32m1_t v_sum = __riscv_vfmv_v_f_f32m1(0.0f, vl); - int n = size; - for (; n > 0;) { - vl = __riscv_vsetvl_e32m1(n); - vfloat32m1_t v_src = __riscv_vle32_v_f32m1(src, vl); - v_sum = __riscv_vfadd_vv_f32m1(v_sum, v_src, vl); - n -= vl; - src += vl; + size_t vlmax = __riscv_vsetvlmax_e32m8(); + vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax); + size_t i = 0; + while (i < size) { + size_t vl = __riscv_vsetvl_e32m8(size - i); + vfloat32m8_t vs = __riscv_vle32_v_f32m8(src + i, vl); + acc = __riscv_vfadd_vv_f32m8_tu(acc, acc, vs, vl); + i += vl; } - vl = __riscv_vsetvlmax_e32m1(); - vfloat32m1_t v_total = __riscv_vfredusum_vs_f32m1_f32m1(v_sum, __riscv_vfmv_s_f_f32m1(0.0f, vl), vl); - float sum = __riscv_vfmv_f_s_f32m1_f32(v_total); - *dst = sum; + vfloat32m1_t sum = __riscv_vfmv_s_f_f32m1(0.0f, 1); + sum = __riscv_vfredusum_vs_f32m8_f32m1(acc, sum, vlmax); + *dst = __riscv_vfmv_f_s_f32m1_f32(sum); } diff --git a/source/backend/metal/ConvSimdGroupShader.hpp b/source/backend/metal/ConvSimdGroupShader.hpp index e55e09c2a7..b82a9c23ac 100644 --- a/source/backend/metal/ConvSimdGroupShader.hpp +++ b/source/backend/metal/ConvSimdGroupShader.hpp @@ -3206,6 +3206,7 @@ template [[host_name("conv1x1_gemv_g4m12_wquant_sg")]] kernel kernel_type_t conv template [[host_name("conv1x1_gemv_g4m13_wquant_sg")]] kernel kernel_type_t conv1x1_gemv_g4mx_wquant_sg<13>; template [[host_name("conv1x1_gemv_g4m14_wquant_sg")]] kernel kernel_type_t conv1x1_gemv_g4mx_wquant_sg<14>; template [[host_name("conv1x1_gemv_g4m15_wquant_sg")]] kernel kernel_type_t conv1x1_gemv_g4mx_wquant_sg<15>; +template [[host_name("conv1x1_gemv_g4m16_wquant_sg")]] kernel kernel_type_t conv1x1_gemv_g4mx_wquant_sg<16>; // Fused weight+scale decode GEMV kernel: scale/bias is stored inline before each weight block // in a single contiguous buffer, eliminating separate scale buffer access. diff --git a/source/backend/opencl/execution/buffer/AttentionBufExecution.cpp b/source/backend/opencl/execution/buffer/AttentionBufExecution.cpp index 273b5d0a24..e59057c639 100644 --- a/source/backend/opencl/execution/buffer/AttentionBufExecution.cpp +++ b/source/backend/opencl/execution/buffer/AttentionBufExecution.cpp @@ -330,14 +330,15 @@ bool KVCacheCLManager::reallocKVCache(const KVMeta* meta, int seqlen, bool isExe mPastLength = start; return true; } - - size_t pastkvSize = mKvNumHead * UP_DIV(mMaxLength, 4) * mHeadDim * 4 * mByte; + size_t curMaxlen = ROUND_UP(mMaxLength, 4); + size_t pastkvSize = mKvNumHead * UP_DIV(curMaxlen, 4) * mHeadDim * 4 * mByte; char* keyPtr = (char*)mOpenCLBackend->getOpenCLRuntime()->commandQueue().enqueueMapBuffer( *mPastKey.get(), true, CL_MAP_READ | CL_MAP_WRITE, 0, pastkvSize, nullptr, nullptr, &res); char* valuePtr = (char*)mOpenCLBackend->getOpenCLRuntime()->commandQueue().enqueueMapBuffer( *mPastValue.get(), true, CL_MAP_READ | CL_MAP_WRITE, 0, pastkvSize, nullptr, nullptr, &res); // TODO: need to ensure reserve info is sorted + auto copyDstIndex = start; for (int n = 0; n < meta->n_reserve; ++n) { auto begin = meta->reserve[2 * n]; auto length = meta->reserve[2 * n + 1]; @@ -345,22 +346,21 @@ bool KVCacheCLManager::reallocKVCache(const KVMeta* meta, int seqlen, bool isExe // past_value : [mKvNumHead, mMaxLength, mHeadDim] auto copySrcIndex = start + begin; - auto copyDstIndex = start; for (int i = 0; i < mKvNumHead * mHeadDim; i++) { - ::memcpy(keyPtr + (i * mMaxLength + copyDstIndex) * mByte, - keyPtr + (i * mMaxLength + copySrcIndex) * mByte, length * mByte); + ::memmove(keyPtr + (i * curMaxlen + copyDstIndex) * mByte, + keyPtr + (i * curMaxlen + copySrcIndex) * mByte, length * mByte); } for (int i = 0; i < mKvNumHead; i++) { for (int j = 0; j < length; j++) { - ::memcpy(valuePtr + (i * mMaxLength + copyDstIndex + j) * mHeadDim * mByte, - valuePtr + (i * mMaxLength + copySrcIndex + j) * mHeadDim * mByte, mHeadDim * mByte); + ::memmove(valuePtr + (i * curMaxlen + copyDstIndex + j) * mHeadDim * mByte, + valuePtr + (i * curMaxlen + copySrcIndex + j) * mHeadDim * mByte, mHeadDim * mByte); } } - start += length; + copyDstIndex += length; } mOpenCLBackend->getOpenCLRuntime()->commandQueue().enqueueUnmapMemObject(*mPastKey.get(), keyPtr); mOpenCLBackend->getOpenCLRuntime()->commandQueue().enqueueUnmapMemObject(*mPastValue.get(), valuePtr); - mPastLength = (int)start; + mPastLength = (int)copyDstIndex; } return true; } diff --git a/source/core/Interpreter.cpp b/source/core/Interpreter.cpp index b5d6a686fe..4f816f4571 100644 --- a/source/core/Interpreter.cpp +++ b/source/core/Interpreter.cpp @@ -81,6 +81,7 @@ static Content* loadModelFile(const char* file) { auto net = new Content; bool success = loader->merge(net->buffer); if (!success) { + delete net; return nullptr; } loader.reset(); @@ -105,7 +106,8 @@ Interpreter* Interpreter::createFromBuffer(const void* buffer, size_t size) { auto net = new Content; net->buffer.reset((int)size); if (nullptr == net->buffer.get()) { - MNN_ERROR("Memory not enought!\n"); + MNN_ERROR("Memory not enough!\n"); + delete net; return nullptr; } ::memcpy(net->buffer.get(), buffer, size); @@ -120,6 +122,7 @@ Interpreter* Interpreter::createFromBufferInternal(Content* net, bool enforceAut } auto valid = OpCommonUtils::checkNet(net->buffer.get(), net->buffer.size()); if (!valid) { + delete net; return nullptr; } net->net = GetNet(net->buffer.get()); diff --git a/test/cv/ImageProcessTest.cpp b/test/cv/ImageProcessTest.cpp index f7d8f666dc..f573d1071d 100644 --- a/test/cv/ImageProcessTest.cpp +++ b/test/cv/ImageProcessTest.cpp @@ -1272,3 +1272,99 @@ class ImageProcessSpeed: public MNNTestCase { } }; // MNNTestSuiteRegister(ImageProcessSpeed, "cv/image_process/speed"); + +// ========== Test: Stride Mismatch ========== +class StrideMismatchTest : public MNNTestCase { +public: + virtual bool run(int precision) { + const int W = 5, H = 5; // Non-power-of-2 to test stride alignment + const int srcChannels = 3; + const int dstChannels = 4; + // Use a wider stride (padded rows) + const int srcStride = W * srcChannels + 4; // Extra padding + std::vector src(H * srcStride, 0); + + // Fill valid pixel data + for (int y = 0; y < H; ++y) { + for (int x = 0; x < W; ++x) { + for (int c = 0; c < srcChannels; ++c) { + src[y * srcStride + x * srcChannels + c] = + static_cast((y * 31 + x * 17 + c * 7) % 256); + } + } + } + + // Use RGB->RGBA conversion to avoid the identity optimization path + // that skips the sampler/blitter pipeline entirely. + ImageProcess::Config config; + config.sourceFormat = RGB; + config.destFormat = RGBA; + config.filterType = MNN::CV::Filter::NEAREST; + config.wrap = CLAMP_TO_EDGE; + + std::unique_ptr process(ImageProcess::create(config)); + MNNTEST_ASSERT(process.get() != nullptr); + + Matrix tr; + process->setMatrix(tr); + + std::vector dst(W * H * dstChannels, 0); + // Pass explicit stride for source; use default for output + process->convert(src.data(), W, H, srcStride, dst.data(), W, H, dstChannels, 0, halide_type_of()); + + // Verify RGB channels match despite stride mismatch, alpha should be 255 + for (int y = 0; y < H; ++y) { + for (int x = 0; x < W; ++x) { + for (int c = 0; c < srcChannels; ++c) { + uint8_t expected = src[y * srcStride + x * srcChannels + c]; + uint8_t actual = dst[(y * W + x) * dstChannels + c]; + MNNTEST_ASSERT(expected == actual); + } + // Alpha channel should be filled (typically 255) + uint8_t alpha = dst[(y * W + x) * dstChannels + 3]; + MNNTEST_ASSERT(alpha == 255); + } + } + return true; + } +}; +MNNTestSuiteRegister(StrideMismatchTest, "cv/image_process/stride_mismatch"); + +// ========== Test: Single-pixel Image ========== +class SinglePixelResizeTest : public MNNTestCase { +public: + virtual bool run(int precision) { + const int channels = 4; + uint8_t src[4] = {100, 150, 200, 255}; + + ImageProcess::Config config; + config.sourceFormat = RGBA; + config.destFormat = RGBA; + config.filterType = MNN::CV::Filter::BILINEAR; + config.wrap = CLAMP_TO_EDGE; + + std::unique_ptr process(ImageProcess::create(config)); + MNNTEST_ASSERT(process.get() != nullptr); + + // Resize 1x1 -> 4x4: all output pixels should equal the source pixel + const int dstW = 4, dstH = 4; + Matrix tr; + float fx = 1.0f / dstW; + float fy = 1.0f / dstH; + tr.postScale(fx, fy); + tr.postTranslate(0.5f * (fx - 1), 0.5f * (fy - 1)); + process->setMatrix(tr); + + std::vector dst(dstW * dstH * channels); + process->convert(src, 1, 1, 0, dst.data(), dstW, dstH, channels, 0, halide_type_of()); + + for (int i = 0; i < dstW * dstH; ++i) { + for (int c = 0; c < channels; ++c) { + // All pixels should be the same as the single source pixel + MNNTEST_ASSERT(std::abs((int)dst[i * channels + c] - (int)src[c]) <= 1); + } + } + return true; + } +}; +MNNTestSuiteRegister(SinglePixelResizeTest, "cv/image_process/single_pixel_resize"); diff --git a/test/expr/LoadMapInputTest.cpp b/test/expr/LoadMapInputTest.cpp new file mode 100644 index 0000000000..08fa6ab802 --- /dev/null +++ b/test/expr/LoadMapInputTest.cpp @@ -0,0 +1,62 @@ +// +// LoadMapInputTest.cpp +// MNNTests +// +// Created by MNN on 2026/08/07. +// Copyright © 2018, Alibaba Group Holding Limited +// + +#include +#include +#include +#include +#include "MNNTestSuite.h" + +using namespace MNN::Express; + +// Regression test for #4731: Variable::loadMap input tensor lost its host buffer, +// making writeMap() return NULL and downstream format conversion crash with a +// null source pointer (e.g. expressDemo SIGSEGV on NCHW-input models). +class LoadMapInputWriteMapTest : public MNNTestCase { +public: + virtual bool run(int precision) override { + // Build a tiny NCHW-input model (Conv3x3 + ReLU) and save it. + auto x = _Input({1, 3, 8, 8}, NCHW); + std::vector weight(4 * 3 * 3 * 3, 0.1f); + std::vector bias(4, 0.01f); + auto w = _Const(weight.data(), {4, 3, 3, 3}, NCHW); + auto b = _Const(bias.data(), {4}, NCHW); + auto y = _Relu(_Conv(w, b, x)); + Variable::save({y}, "regression_4731.mnn"); + + // Load the model and check the input VARP is writable. + auto varMap = Variable::loadMap("regression_4731.mnn"); + auto io = Variable::getInputAndOutput(varMap); + if (io.first.empty() || io.second.empty()) { + MNN_PRINT("LoadMapInputTest: no input/output found\n"); + return false; + } + auto input = io.first.begin()->second; + auto ptr = input->writeMap(); + if (nullptr == ptr) { + // Before the fix this was NULL (input tensor host was dropped by + // Tensor::clone in Variable::load), causing SIGSEGV downstream. + MNN_PRINT("LoadMapInputTest: writeMap returned NULL (bug #4731)\n"); + return false; + } + auto inInfo = input->getInfo(); + int size = 1; + for (auto d : inInfo->dim) { + size *= d; + } + for (int i = 0; i < size; ++i) { + ptr[i] = 0.5f; + } + // Forward must not crash (compute succeeds; output reading is a separate + // follow-up concern, see issue #4731). + auto output = io.second.begin()->second; + (void)output->readMap(); + return true; + } +}; +MNNTestSuiteRegister(LoadMapInputWriteMapTest, "expr/LoadMapInputWriteMap"); diff --git a/transformers/llm/engine/src/llm.cpp b/transformers/llm/engine/src/llm.cpp index fa7d354dde..3c524aabc8 100644 --- a/transformers/llm/engine/src/llm.cpp +++ b/transformers/llm/engine/src/llm.cpp @@ -427,7 +427,7 @@ bool Llm::load() { } if (mConfig->is_mrope()) { - mPositionIdsVarVec[i] = _Input({3, index}, NCHW, halide_type_of()); + mPositionIdsVarVec[i] = _Input({mConfig->mrope_axes(), index}, NCHW, halide_type_of()); } else { mPositionIdsVarVec[i] = _Input({1, index}, NCHW, halide_type_of()); } @@ -1615,8 +1615,9 @@ VARP Llm::gen_position_ids(int seq_len) { auto ptr = mPositionIdsVarVec[0]->writeMap(); ptr[0] = is_glm2 ? mContext->gen_seq_len : mContext->all_seq_len; if (mConfig->is_mrope()) { - ptr[1] = ptr[0]; - ptr[2] = ptr[0]; + for (int axis = 1; axis < mConfig->mrope_axes(); axis++) { + ptr[axis] = ptr[0]; + } } return mPositionIdsVarVec[0]; } @@ -1625,16 +1626,21 @@ VARP Llm::gen_position_ids(int seq_len) { for (int i = 0; i < seq_len; i++) { ptr[i] = i + mContext->all_seq_len; } + if (mConfig->is_mrope()) { + for (int axis = 1; axis < mConfig->mrope_axes(); axis++) { + ::memcpy(ptr + axis * seq_len, ptr, seq_len * sizeof(int)); + } + } return mPositionIdsVarVec[1]; } if (mConfig->is_mrope()) { - positionIds = _Input({3, seq_len}, NCHW, halide_type_of()); + positionIds = _Input({mConfig->mrope_axes(), seq_len}, NCHW, halide_type_of()); auto ptr = positionIds->writeMap(); - for (int i = 0; i < seq_len; i++) { - ptr[0 * seq_len + i] = i + mContext->all_seq_len; - ptr[1 * seq_len + i] = i + mContext->all_seq_len; - ptr[2 * seq_len + i] = i + mContext->all_seq_len; + for (int axis = 0; axis < mConfig->mrope_axes(); axis++) { + for (int i = 0; i < seq_len; i++) { + ptr[axis * seq_len + i] = i + mContext->all_seq_len; + } } return positionIds; } diff --git a/transformers/llm/engine/src/llmconfig.hpp b/transformers/llm/engine/src/llmconfig.hpp index e95c672352..1cfb2ad014 100644 --- a/transformers/llm/engine/src/llmconfig.hpp +++ b/transformers/llm/engine/src/llmconfig.hpp @@ -269,6 +269,8 @@ class LlmConfig { return config_.value("is_mrope", false); } + int mrope_axes() const { return config_.value("mrope_axes", 3); } + bool has_talker() const { return config_.value("has_talker", false); } @@ -634,4 +636,4 @@ class LlmConfig { } // Transformer } // MNN -#endif \ No newline at end of file +#endif diff --git a/transformers/llm/engine/src/omni.cpp b/transformers/llm/engine/src/omni.cpp index e800bb522e..db8ea2044a 100644 --- a/transformers/llm/engine/src/omni.cpp +++ b/transformers/llm/engine/src/omni.cpp @@ -11,6 +11,7 @@ #endif #include #include +#include #include #include #include @@ -510,6 +511,91 @@ std::vector Omni::qwen2VisionProcess(VARP image) { return imgIds; } +std::vector Omni::hunyuanVisionProcess(VARP image) { + MNN::Express::ExecutorScope s(mExecutor); + int patchSize = mConfig->config_.value("hunyuan_patch_size", 16); + int mergeSize = mConfig->config_.value("hunyuan_spatial_merge_size", 2); + int temporalPatchSize = mConfig->config_.value("hunyuan_temporal_patch_size", 1); + if (patchSize <= 0 || mergeSize <= 0 || temporalPatchSize <= 0) { + MNN_ERROR("Invalid Hunyuan vision config: patch=%d merge=%d temporal=%d\n", patchSize, mergeSize, + temporalPatchSize); + return std::vector(0); + } + if (temporalPatchSize != 1) { + MNN_ERROR("Hunyuan temporal_patch_size=%d is not supported by Omni image preprocessing\n", temporalPatchSize); + return std::vector(0); + } + if (!mVisionSizeOverridden) { + auto imageInfo = image->getInfo(); + if (imageInfo != nullptr && imageInfo->dim.size() >= 2) { + auto dims = imageInfo->dim; + int imageHeight = dims[0]; + int imageWidth = dims[1]; + if (dims.size() >= 3 && dims[dims.size() - 1] <= 4) { + imageHeight = dims[dims.size() - 3]; + imageWidth = dims[dims.size() - 2]; + } + if (imageHeight > 0 && imageWidth > 0) { + mVisionHeight = imageHeight; + mVisionWidth = imageWidth; + } + } + } + const int factor = patchSize * mergeSize; + int minPixels = mConfig->config_.value("image_min_pixels", mVisionHeight * mVisionWidth); + int maxPixels = mConfig->config_.value("image_max_pixels", mVisionMaxSize * mVisionMaxSize); + int resizedHeight = + std::max(factor, static_cast(std::round(static_cast(mVisionHeight) / factor)) * factor); + int resizedWidth = + std::max(factor, static_cast(std::round(static_cast(mVisionWidth) / factor)) * factor); + if (resizedHeight * resizedWidth > maxPixels) { + float beta = std::sqrt(static_cast(mVisionHeight * mVisionWidth) / maxPixels); + resizedHeight = std::max(factor, static_cast(std::floor(mVisionHeight / beta / factor)) * factor); + resizedWidth = std::max(factor, static_cast(std::floor(mVisionWidth / beta / factor)) * factor); + } else if (resizedHeight * resizedWidth < minPixels) { + float beta = std::sqrt(static_cast(minPixels) / (mVisionHeight * mVisionWidth)); + resizedHeight = std::max(factor, static_cast(std::ceil(mVisionHeight * beta / factor)) * factor); + resizedWidth = std::max(factor, static_cast(std::ceil(mVisionWidth * beta / factor)) * factor); + } + mVisionHeight = resizedHeight; + mVisionWidth = resizedWidth; + image = MNN::CV::resize(image, {mVisionWidth, mVisionHeight}, 0, 0, MNN::CV::INTER_CUBIC, MNN::CV::COLOR_BGR2RGB, + mVisionMean, mVisionNorm); + image = Express::_Unsqueeze(image, {0}); + image = Express::_Convert(image, NCHW); + int gridH = mVisionHeight / patchSize; + int gridW = mVisionWidth / patchSize; + auto patches = Express::_Reshape( + image, {1, 3, gridH / mergeSize, mergeSize, patchSize, gridW / mergeSize, mergeSize, patchSize}); + patches = Express::_Permute(patches, {0, 2, 3, 5, 6, 1, 4, 7}); + patches = Express::_Reshape(patches, {gridH * gridW, 3 * temporalPatchSize * patchSize * patchSize}); + auto imageGridThw = Express::_Input({1, 3}, NCHW, halide_type_of()); + auto gridPtr = imageGridThw->writeMap(); + gridPtr[0] = 1; + gridPtr[1] = gridH; + gridPtr[2] = gridW; + auto outputs = mVisionModule->onForward({patches, imageGridThw}); + if (outputs.empty() || outputs[0] == nullptr || outputs[0]->getInfo() == nullptr) { + MNN_ERROR("Hunyuan vision forward failed: resized=%dx%d grid=%dx%d patch=%d merge=%d\n", mVisionHeight, + mVisionWidth, gridH, gridW, patchSize, mergeSize); + return std::vector(0); + } + auto imageEmbedding = outputs[0]; + int visionLen = imageEmbedding->getInfo()->dim[0]; + int gridTokens = (gridH / mergeSize) * (gridW / mergeSize + 1); + int extraTokens = visionLen - gridTokens; + if (extraTokens != 0 && extraTokens != 2) { + MNN_ERROR("Hunyuan image token count mismatch: tokens=%d grid=%d\n", visionLen, gridTokens); + return std::vector(0); + } + mVisionEmbeddings.push_back(imageEmbedding); + addPositionIds(visionLen, gridH / mergeSize, gridW / mergeSize); + std::vector imgIds(visionLen, mVisionPad); + imgIds.insert(imgIds.begin(), mVisionStart); + imgIds.push_back(mVisionEnd); + return imgIds; +} + std::vector Omni::smolvlmVisionProcess(VARP image) { MNN::Express::ExecutorScope s(mExecutor); // SmolVLM / LFM2-VL: compute visionLen from global image forward @@ -804,13 +890,21 @@ std::vector Omni::visionProcess(VARP image) { #ifdef LLM_SUPPORT_VISION if (image == nullptr) { MNN_PRINT("Omni Can't open image\n"); + mVisionSizeOverridden = false; return std::vector(0); } Timer _t; std::vector imgIds; const auto inputNames = mVisionModule->getInfo()->inputNames; + const auto visionType = mConfig->config_.value("vision_type", ""); if (inputNames.size() >= 3 && inputNames[0] == "patches") { imgIds = qwen2VisionProcess(image); + } else if (visionType == "hunyuan_vl") { + if (inputNames.size() == 2 && inputNames[0] == "pixel_values" && inputNames[1] == "image_grid_thw") { + imgIds = hunyuanVisionProcess(image); + } else { + MNN_ERROR("Hunyuan vision expects inputs pixel_values,image_grid_thw\n"); + } } else if (inputNames[0] == "pixel_values") { if (inputNames.size() == 1) { imgIds = smolvlmVisionProcess(image); @@ -830,6 +924,7 @@ std::vector Omni::visionProcess(VARP image) { } mContext->vision_us += _t.durationInUs(); mContext->pixels_mp += (mVisionWidth / 1000.0f) * (mVisionHeight / 1000.0f); + mVisionSizeOverridden = false; // set vision number for image idx mVisionNum += 1; return imgIds; @@ -965,7 +1060,12 @@ std::vector Omni::multimodeProcess(const std::string& mode, std::string inf std::stringstream hw_ss(match.str(1)); char comma; - hw_ss >> mVisionHeight >> comma >> mVisionWidth; + int parsedHeight = 0, parsedWidth = 0; + if (hw_ss >> parsedHeight >> comma >> parsedWidth && parsedHeight > 0 && parsedWidth > 0) { + mVisionHeight = parsedHeight; + mVisionWidth = parsedWidth; + mVisionSizeOverridden = true; + } currentPosition = matchPosition + match.length(); } if (currentPosition < info.length()) { @@ -1011,6 +1111,29 @@ std::vector Omni::multimodeProcess(const std::string& mode, std::string inf } void Omni::addPositionIds(int t, int h, int w) { + if (mConfig->config_.value("vision_type", "") == "hunyuan_vl" && h >= 0 && w >= 0) { + int cur_idx = mPositionIds.mT.empty() ? 0 : mPositionIds.mT.back() + 1; + int gridTokens = h * (w + 1); + int extraTokens = t - gridTokens; + if (extraTokens != 0 && extraTokens != 2) { + MNN_ERROR("Hunyuan image token count mismatch: tokens=%d grid=%d\n", t, gridTokens); + return; + } + mPositionIds.push_back(cur_idx++); + if (extraTokens == 2) { + mPositionIds.push_back(cur_idx++); + } + for (int h_i = 0; h_i < h; h_i++) { + for (int w_i = 0; w_i <= w; w_i++) { + mPositionIds.push_back(cur_idx++, w_i, h_i, mVisionNum); + } + } + if (extraTokens == 2) { + mPositionIds.push_back(cur_idx++); + } + mPositionIds.push_back(cur_idx++); + return; + } int cur_idx = mPositionIds.currentIdx(); if (h < 0 && w < 0) { // text position ids for (int i = 0; i < t; i++) { @@ -1040,6 +1163,7 @@ std::vector Omni::tokenizer_encode(const MultimodalPrompt& multimodal_input std::smatch match; std::vector ids{}; mPositionIds.clear(); + mVisionNum = 0; while (std::regex_search(searchStart, prompt.cend(), match, multimode_regex)) { auto txt_ids = mTokenizer->encode(match.prefix().str()); @@ -1079,6 +1203,7 @@ std::vector Omni::processImageContent(const std::string& content, const std if (it->second.height > 0 && it->second.width > 0) { mVisionHeight = it->second.height; mVisionWidth = it->second.width; + mVisionSizeOverridden = true; } // MNN_PRINT("processImageContent: using placeholder '%s' with size %dx%d", content.c_str(), mVisionWidth, mVisionHeight); return visionProcess(it->second.image_data); @@ -1237,26 +1362,42 @@ VARP Omni::gen_position_ids(int seq_len) { return Llm::gen_position_ids(seq_len); } // mrope - if (needNewVar(positionIds, 1, seq_len)) { - positionIds = _Input({3, seq_len}, NCHW, halide_type_of()); + int axes = mConfig->mrope_axes(); + if (positionIds == nullptr || positionIds->getInfo()->dim[0] != axes || needNewVar(positionIds, 1, seq_len)) { + positionIds = _Input({axes, seq_len}, NCHW, halide_type_of()); } auto ptr = positionIds->writeMap(); if (mContext->gen_seq_len > 0) { - for (int i=0; igen_seq_len + mPositionIds.back() + i; auto pos = mContext->all_seq_len + i; - ptr[i + 0] = pos; - ptr[i + seq_len] = pos; - ptr[i + seq_len * 2] = pos; + for (int axis = 0; axis < axes; axis++) { + ptr[i + seq_len * axis] = pos; + } } } else { + bool hunyuan = mConfig->config_.value("vision_type", "") == "hunyuan_vl"; + auto axisValue = [this](int axis, int i) { + const std::vector* values = nullptr; + if (axis == 0) { + values = &mPositionIds.mT; + } else if (axis == 1) { + values = &mPositionIds.mH; + } else if (axis == 2) { + values = &mPositionIds.mW; + } else if (axis == 3) { + values = &mPositionIds.mX; + } + if (values != nullptr && i < static_cast(values->size())) { + return (*values)[i]; + } + return i; + }; for (int i = 0; i < seq_len; i++) { - int mT_val = i < mPositionIds.mT.size() ? mPositionIds.mT[i] : i; - int mH_val = i < mPositionIds.mH.size() ? mPositionIds.mH[i] : i; - int mW_val = i < mPositionIds.mW.size() ? mPositionIds.mW[i] : i; - ptr[i] = mT_val + mContext->all_seq_len; - ptr[i + seq_len] = mH_val + mContext->all_seq_len; - ptr[i + seq_len * 2] = mW_val + mContext->all_seq_len; + for (int axis = 0; axis < axes; axis++) { + int offset = (hunyuan && axis > 0) ? 0 : mContext->all_seq_len; + ptr[i + seq_len * axis] = axisValue(axis, i) + offset; + } } if (mTalker) { mTalker->setPostionIds(mPositionIds); diff --git a/transformers/llm/engine/src/omni.hpp b/transformers/llm/engine/src/omni.hpp index 5ab8c9fdbe..7362b35f44 100644 --- a/transformers/llm/engine/src/omni.hpp +++ b/transformers/llm/engine/src/omni.hpp @@ -27,6 +27,7 @@ class MropeInfo { mT = info.mT; mH = info.mH; mW = info.mW; + mX = info.mX; } int back() { if (mW.empty()) { @@ -40,10 +41,12 @@ class MropeInfo { } return back() + 1; } - void push_back(int t, int h, int w) { + void push_back(int t, int h, int w) { push_back(t, h, w, w); } + void push_back(int t, int h, int w, int x) { mT.push_back(t); mH.push_back(h); mW.push_back(w); + mX.push_back(x); } void push_back(int t) { push_back(t, t, t); @@ -56,8 +59,9 @@ class MropeInfo { mT.clear(); mH.clear(); mW.clear(); + mX.clear(); } - std::vector mT, mH, mW; + std::vector mT, mH, mW, mX; }; struct WavChunk { @@ -165,6 +169,7 @@ class Omni : public Llm { std::vector qwen2VisionProcess(VARP image); std::vector smolvlmVisionProcess(VARP image); std::vector minicpmVisionProcess(VARP image); + std::vector hunyuanVisionProcess(VARP image); std::vector gemma4VisionProcess(VARP image); private: int mVisionHeight = 448, mVisionWidth = 448, mVisionStart = 151857, @@ -173,6 +178,7 @@ class Omni : public Llm { int mVisionGlobal = 49152; int mVisionSizeUnit = 1, mVisionMaxSize = 2048; int mVisionNum = 0; + bool mVisionSizeOverridden = false; std::vector mVisionMean{122.7709383, 116.7460125, 104.09373615}; std::vector mVisionNorm{0.01459843, 0.01500777, 0.01422007}; std::vector multimodeProcess(const std::string& mode, std::string info); diff --git a/transformers/llm/engine/tools/generateLlmIO.cpp b/transformers/llm/engine/tools/generateLlmIO.cpp index 39f729ac9d..3daf4a994d 100644 --- a/transformers/llm/engine/tools/generateLlmIO.cpp +++ b/transformers/llm/engine/tools/generateLlmIO.cpp @@ -102,26 +102,63 @@ static void createInputsForEmbedding(int seqLen, int hiddenSize, const std::stri inputs.push_back(positionIds); } -static bool isEmbeddingModel(const rapidjson::Document& doc) { - if (doc.HasMember("output_names") && doc["output_names"].IsArray()) { - for (auto iter = doc["output_names"].Begin(); iter != doc["output_names"].End(); ++iter) { - if (iter->IsString() && std::string(iter->GetString()) == "sentence_embeddings") { - return true; - } +static bool hasStringValue(const rapidjson::Document& doc, const char* key, const std::string& value) { + if (!doc.HasMember(key)) { + return false; + } + auto& member = doc[key]; + if (member.IsString()) { + return member.GetString() == value; + } + if (!member.IsArray()) { + return false; + } + for (auto iter = member.Begin(); iter != member.End(); ++iter) { + if (iter->IsString() && iter->GetString() == value) { + return true; + } + } + return false; +} + +static bool isEmbeddingModel(const rapidjson::Document& doc, const std::string& modelDir) { + if (doc.HasMember("is_embedding") && doc["is_embedding"].IsBool() && doc["is_embedding"].GetBool()) { + return true; + } + if (hasStringValue(doc, "output_names", "sentence_embeddings")) { + return true; + } + if (doc.HasMember("llm_model") && doc["llm_model"].IsString()) { + auto modelName = std::string(doc["llm_model"].GetString()); + if (modelName.find("embedding.mnn") != std::string::npos) { + return true; + } + } + if (doc.HasMember("embedding_model") && doc["embedding_model"].IsString()) { + auto modelName = std::string(doc["embedding_model"].GetString()); + if (modelName.find("embedding.mnn") != std::string::npos) { + return true; } } + if (MNNFileExist(MNNFilePathConcat(modelDir, "embedding.mnn").c_str())) { + return true; + } auto modelType = std::string(doc.HasMember("model_type") && doc["model_type"].IsString() ? doc["model_type"].GetString() : ""); - if (modelType == "bert" || modelType == "new" || modelType == "qwen3") { + if (modelType == "bert" || modelType == "new") { + return true; + } + if (modelType == "qwen3" && !doc.HasMember("layer_nums") && !doc.HasMember("attention_type") && !doc.HasMember("is_mrope")) { return true; } return false; } -static bool generateForModel(const std::string& modelPath, const std::string& outputDir, const std::string& jsonPath, int blockSize) { +static bool generateForModel(const std::string& modelDir, const std::string& outputDir, const std::string& jsonPath, int blockSize) { std::shared_ptr net; std::vector inputNames; std::vector outputNames; bool isEmbedding = false; + std::string modelPath; int hiddenSize; std::string attentionMaskType; @@ -152,8 +189,17 @@ static bool generateForModel(const std::string& modelPath, const std::string& ou } attentionMaskType = doc["attention_mask"].GetString(); - isEmbedding = isEmbeddingModel(doc); + isEmbedding = isEmbeddingModel(doc, modelDir); + } + + modelPath = MNNFilePathConcat(modelDir, isEmbedding ? "embedding.mnn" : "llm.mnn"); + if (isEmbedding && !MNNFileExist(modelPath.c_str())) { + auto llmModelPath = MNNFilePathConcat(modelDir, "llm.mnn"); + if (MNNFileExist(llmModelPath.c_str())) { + modelPath = llmModelPath; + } } + FUNC_PRINT_ALL(modelPath.c_str(), s); MNN::ScheduleConfig config; std::shared_ptr rtmgr(MNN::Express::Executor::RuntimeManager::createRuntimeManager(config)); @@ -229,9 +275,8 @@ int main(int argc, char* argv[]) { } FUNC_PRINT(blockSize); - std::string modelPath = std::string(argv[1]) + "/llm.mnn"; - std::string llmConfigPath = std::string(argv[1]) + "/llm_config.json"; - FUNC_PRINT_ALL(modelPath.c_str(), s); + std::string modelDir = argv[1]; + std::string llmConfigPath = MNNFilePathConcat(modelDir, "llm_config.json"); FUNC_PRINT_ALL(llmConfigPath.c_str(), s); std::string outputDir = argv[2]; @@ -239,7 +284,7 @@ int main(int argc, char* argv[]) { MNN_PRINT("Failed to create dir %s.\n", outputDir.c_str()); } - if (!generateForModel(modelPath, outputDir, llmConfigPath, blockSize)) { + if (!generateForModel(modelDir, outputDir, llmConfigPath, blockSize)) { return 1; } diff --git a/transformers/llm/export/llmexport.py b/transformers/llm/export/llmexport.py index 2d4e2d00cf..b9a4e25420 100644 --- a/transformers/llm/export/llmexport.py +++ b/transformers/llm/export/llmexport.py @@ -86,6 +86,8 @@ def visit_module(module): 'attention_type': self.config.attention_type, 'is_mrope': self.model.rotary.is_mrope } + if self.model.rotary.is_mrope: + self.llm_config['mrope_axes'] = self.model.rotary.mrope_axes self.llm_config.update(self.model.get_config()) # Attention scaling (gemma4 uses 1.0 instead of 1/sqrt(head_dim)) if hasattr(self.model, 'blocks') and len(self.model.blocks) > 0: @@ -117,6 +119,13 @@ def visit_module(module): 'chat_template': "[gMASK]{% for message in messages %}{% if message.role == \"user\" %}<|user|>\n{{ message.content }}{% elif message.role == \"assistant\" %}<|assistant|>\n{{ message.content }}{% elif message.role == \"system\" %}<|system|>\n{{ message.content }}{% endif %}{% endfor %}{% if add_generation_prompt %}<|assistant|>\n{% endif %}", 'eos': '<|endoftext|>' } + # HunyuanVL's HF template uses syntax unsupported by the C++ minja parser. + if self.model_type == 'hunyuan_vl': + self.llm_config['jinja'] = { + 'chat_template': "<|hy_begin▁of▁sentence|>{% for message in messages %}{% if message.role == \"system\" %}{{ message.content }}<|hy_place▁holder▁no▁3|>{% elif message.role == \"user\" %}{{ message.content }}<|hy_User|>{% elif message.role == \"assistant\" %}{{ message.content }}<|hy_Assistant|>{% endif %}{% endfor %}", + 'bos': '<|hy_begin▁of▁sentence|>', + 'eos': '<|hy_Assistant|>' + } # tie word embeddings self.args.tie_word_embeddings = not self.args.seperate_embed and self.model.lm.lm.weight.equal(self.model.embed.embed.weight) @@ -410,6 +419,8 @@ def build_faker(real, name): for i in range(len(self.model.blocks)): # different kv cache shape in different layers # if isinstance(self.config.num_attention_heads, list): + # Keep custom Attention in the exported LLM graph so runtime decode uses KV cache. + # HunyuanVL still disables MNNConvert transformerFuse separately. self.model.blocks[i].self_attn.export_fused_attn = True is_moe = hasattr(self.model.blocks[i].mlp, 'is_moe') and self.model.blocks[i].mlp.is_moe if is_moe: @@ -672,7 +683,8 @@ def export_language(self): if self.args.onnx_slim: self.slim_onnx(onnx_model) if self.mnn_converter: - tie_embeddings_info = MNNConverter(self, self.unloaded_ops).export(onnx_model) + fuse_transformer = self.model_type != 'hunyuan_vl' + tie_embeddings_info = MNNConverter(self, self.unloaded_ops).export(onnx_model, transformer_fuse=fuse_transformer) if tie_embeddings_info is not None: self.llm_config['tie_embeddings'] = tie_embeddings_info else: @@ -713,6 +725,7 @@ def export_tokenizer(self): class EmbeddingExporter(LlmExporter): def __init__(self, args): super().__init__(args) + self.dst_name = 'embedding' def response(self, query): self.model.eval() diff --git a/transformers/llm/export/npu/generate_llm_qnn.py b/transformers/llm/export/npu/generate_llm_qnn.py index 9cf993e83c..4f9fd20250 100644 --- a/transformers/llm/export/npu/generate_llm_qnn.py +++ b/transformers/llm/export/npu/generate_llm_qnn.py @@ -21,66 +21,74 @@ def makeIO(args, model_name, inputjson, external_file = None): process.wait() return process.returncode -def makeIOJson(args, seq_len, hidden_size, mask_type): +def is_embedding_model(config_data, model_dir): + if config_data.get("is_embedding", False) is True: + return True + + output_names = config_data.get("output_names", []) + if isinstance(output_names, str): + output_names = [output_names] + if "sentence_embeddings" in output_names: + return True + + model_name = config_data.get("llm_model", config_data.get("embedding_model", "")) + if os.path.basename(model_name) == "embedding.mnn": + return True + + if os.path.exists(os.path.join(model_dir, "embedding.mnn")): + return True + + model_type = config_data.get("model_type", "") + if model_type in ("bert", "new"): + return True + if model_type == "qwen3" and not any(key in config_data for key in ("layer_nums", "attention_type", "is_mrope")): + return True + return False + +def makeIOJson(args, seq_len, hidden_size, mask_type, is_embedding=False): + def model_inputs(current_seq_len, logits_index=None): + inputs = [ + { + "name": "input_ids", + "shape": [current_seq_len, 1, hidden_size] + }, + { + "name": "attention_mask", + "shape": [1, 1, current_seq_len, current_seq_len], + "type": mask_type + }, + { + "name": "position_ids", + "shape": [1, current_seq_len], + "type": "int" + } + ] + if logits_index is not None: + inputs.append({ + "name": "logits_index", + "shape": [1], + "type": "int", + "value": logits_index + }) + return inputs + config = { "configs": [ { - "inputs": [ - { - "name": "input_ids", - "shape": [seq_len, 1, hidden_size] - }, - { - "name": "attention_mask", - "shape": [1, 1, seq_len, seq_len], - "type": mask_type - }, - { - "name": "position_ids", - "shape": [1, seq_len], - "type": "int" - }, - { - "name": "logits_index", - "shape": [1], - "type": "int", - "value": 0 - } - ], + "inputs": model_inputs(seq_len, None if is_embedding else 0), "outputs": [ - "logits" + "sentence_embeddings" if is_embedding else "logits" ] }, { - "inputs": [ - { - "name": "input_ids", - "shape": [1, 1, hidden_size] - }, - { - "name": "attention_mask", - "shape": [1, 1, 1, 1], - "type": mask_type - }, - { - "name": "position_ids", - "shape": [1, 1], - "type": "int" - }, - { - "name": "logits_index", - "shape": [1], - "type": "int", - "value": -1 - } - ], + "inputs": model_inputs(1, None if is_embedding else -1), "outputs": [ - "logits" + "sentence_embeddings" if is_embedding else "logits" ] } ] } - if "Qwen3.5" in args.model: + if not is_embedding and "Qwen3.5" in args.model: cfg = config["configs"] inputs = cfg[0]["inputs"] for inp in inputs: @@ -95,7 +103,7 @@ def makeIOJson(args, seq_len, hidden_size, mask_type): inp["shape"] = [2, 1, 1, 1, 3] if inp["name"] == "position_ids": inp["shape"] = [3, 1] - if "Qwen" in args.model and "VL" in args.model: + if not is_embedding and "Qwen" in args.model and "VL" in args.model: cfg = config["configs"] inputs = cfg[0]["inputs"] for inp in inputs: @@ -243,7 +251,7 @@ def compile_qnn(args): process.wait() return process.returncode -def output_qnn(args): +def output_qnn(args, model_name=None): if os.path.exists(os.path.join(args.model, 'qnn')): shutil.rmtree(os.path.join(args.model, 'qnn')) shutil.move(os.path.join(args.cache_path, 'qnn'), os.path.join(args.model, 'qnn')) @@ -253,9 +261,11 @@ def output_qnn(args): if os.path.exists(config_path): with open(config_path, 'r', encoding='utf-8') as f: config_npu = json.load(f) - is_visual = args.model_name == "visual.mnn" + model_name = model_name or args.model_name + is_visual = model_name == "visual.mnn" if not is_visual: - config_npu["llm_model"] = "qnn/llm.mnn" + config_npu["llm_model"] = "qnn/" + model_name + config_npu["llm_weight"] = model_name + ".weight" config_npu["chunk_limits"] = [args.chunk_size, 1] else: config_npu["visual_model"] = "qnn/visual.mnn" @@ -283,7 +293,7 @@ def convert_qnn(args, model_name, inputjson, external_file, ids): end = time.time() print("Cost: ", end - sta, ' s') print("Step4: Move result file to ", args.model) - output_qnn(args) + output_qnn(args, model_name) print("End") @@ -323,22 +333,24 @@ def convert_llm(args): os.makedirs(cache, exist_ok=True) hidden_size = 768 mask_type = "int" - config_file_path = os.path.join(os.getcwd(), args.model, 'llm_config.json') + model_dir = os.path.join(os.getcwd(), args.model) + config_file_path = os.path.join(model_dir, 'llm_config.json') with open(config_file_path, 'r', encoding='utf-8') as f: config_data = json.load(f) if "hidden_size" in config_data: hidden_size = config_data["hidden_size"] else: - print(f"Error: 'hidden_size' key not found in {config_file_path}") - return npu_convert + raise KeyError(f"'hidden_size' key not found in {config_file_path}") if "attention_mask" in config_data: mask_type = config_data["attention_mask"] + is_embedding = is_embedding_model(config_data, model_dir) ids = [0, 1] - external_file = os.path.join(os.getcwd(), args.model, 'llm.mnn.weight') - makeIOJson(args, args.chunk_size, hidden_size, mask_type) + model_name = 'embedding.mnn' if is_embedding else 'llm.mnn' + external_file = os.path.join(model_dir, model_name + '.weight') + makeIOJson(args, args.chunk_size, hidden_size, mask_type, is_embedding) inputjson = os.path.join(cache, 'input.json') - convert_qnn(args, 'llm.mnn', inputjson, external_file, ids) + convert_qnn(args, model_name, inputjson, external_file, ids) def convert_input_json(args): cache = os.path.join(os.getcwd(), args.cache_path) diff --git a/transformers/llm/export/utils/custom_op.py b/transformers/llm/export/utils/custom_op.py index 79dcb7e3ab..d9197df214 100644 --- a/transformers/llm/export/utils/custom_op.py +++ b/transformers/llm/export/utils/custom_op.py @@ -47,8 +47,7 @@ def symbolic(g, query, key, value, attention_mask, output_dim, kv_cache, name, l "head_dim_i": head_dim, } from torch.onnx.symbolic_helper import _get_tensor_sizes - out_sizes = _get_tensor_sizes(query) - out_sizes[-1] = output_dim + out_sizes = _get_tensor_sizes(query)[:2] + [output_dim] output_type = query.type().with_sizes(out_sizes) return g.op("LlmExporter::FusedAttention", query, key, value, attention_mask, **kwargs).setType(output_type) diff --git a/transformers/llm/export/utils/dflash.py b/transformers/llm/export/utils/dflash.py index cfe9b62557..740761b411 100644 --- a/transformers/llm/export/utils/dflash.py +++ b/transformers/llm/export/utils/dflash.py @@ -5,7 +5,7 @@ from typing import Optional, Tuple from .transformers import Attention, RMSNorm, Rotary, Embedding -from utils.custom_op import FakeLinear +from utils.custom_op import FakeLinear, FusedAttention from utils.spinner import spinner_run from .torch_utils import onnx_export from transformers.activations import ACT2FN @@ -29,6 +29,13 @@ def __init__(self, config, layer_idx): self.q_norm = RMSNorm(self.head_dim) self.k_norm = RMSNorm(self.head_dim) + self.fused_attn = FusedAttention( + self.num_attention_heads * self.head_dim, + kv_cache=0, + name=f'/dflash_layers.{layer_idx}/self_attn/FusedAttention', + layer_index=-1, + kv_shared_layer_index=-1) + def forward(self, hidden_states, context_hidden, q_cos, q_sin, k_cos, k_sin, attention_mask): """ hidden_states: [1, block_size, hidden_size] (noise) @@ -42,40 +49,26 @@ def forward(self, hidden_states, context_hidden, q_cos, q_sin, k_cos, k_sin, att ctx_len = context_hidden.shape[1] total_len = ctx_len + q_len - # Q from noise only - q = self.q_proj(hidden_states) - q = q.view(bsz, q_len, self.num_attention_heads, self.head_dim) - q = self.q_norm(q).transpose(1, 2) # [1, num_heads, q_len, head_dim] - - # K/V from cat(context, noise) + # Projections + q/k norm in [B, seq, heads, head_dim] layout + q = self.q_norm(self.q_proj(hidden_states).view(bsz, q_len, self.num_attention_heads, self.head_dim)) kv_input = torch.cat([context_hidden, hidden_states], dim=1) # [1, total_len, hidden_size] - k = self.k_proj(kv_input) - v = self.v_proj(kv_input) - k = k.view(bsz, total_len, self.num_key_value_heads, self.head_dim) - k = self.k_norm(k).transpose(1, 2) # [1, num_kv_heads, total_len, head_dim] - v = v.view(bsz, total_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + k = self.k_norm(self.k_proj(kv_input).view(bsz, total_len, self.num_key_value_heads, self.head_dim)) + v = self.v_proj(kv_input).view(bsz, total_len, self.num_key_value_heads, self.head_dim) - # Apply RoPE (pre-computed, no dynamic slicing needed) + # RoPE, then one fused Attention op (K/V un-repeated; the op does GQA internally) q = self._apply_rope(q, q_cos, q_sin) k = self._apply_rope(k, k_cos, k_sin) + attn_output = self.fused_attn(q, k, v, attention_mask) # [1, q_len, num_heads*head_dim] - # GQA repeat - if self.num_key_value_groups > 1: - k = k.repeat_interleave(self.num_key_value_groups, dim=1) - v = v.repeat_interleave(self.num_key_value_groups, dim=1) - - # Attention - attn_weights = torch.matmul(q, k.transpose(-2, -1)) * self.scaling - attn_weights = attn_weights + attention_mask - attn_weights = torch.softmax(attn_weights, dim=-1) - attn_output = torch.matmul(attn_weights, v) - - attn_output = attn_output.transpose(1, 2).reshape(bsz, q_len, -1) + # No-op reshape: FusedAttentionOp.symbolic annotates the output as rank-4 while the runtime tensor is rank-3 + attn_output = attn_output.reshape(bsz, q_len, -1) return self.o_proj(attn_output) @staticmethod def _apply_rope(x, cos, sin): - """Apply rotary position embedding.""" + """RoPE for [B, seq, heads, dim] layout: transpose cos/sin [1,1,seq,dim]->[1,seq,1,dim].""" + cos = cos.transpose(1, 2) + sin = sin.transpose(1, 2) x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] rotated = torch.cat((-x2, x1), dim=-1) diff --git a/transformers/llm/export/utils/mnn_converter.py b/transformers/llm/export/utils/mnn_converter.py index d318ca6318..9a206441f9 100644 --- a/transformers/llm/export/utils/mnn_converter.py +++ b/transformers/llm/export/utils/mnn_converter.py @@ -163,7 +163,10 @@ def export(self, onnx_path, quant_bit = None, quant_block = None, transformer_fu self.mnn2json(self.mnn_model_path, mnn_json) self.rebuild(mnn_json) self.json2mnn(mnn_json, self.mnn_model_path) - self.removeDupOps(self.mnn_model_path) + # HunyuanVL needs the explicit q/k/v reshape boundary before Attention. + # The extra optimize pass can bypass v_proj's post_reshape into Attention. + if self.exporter.model_type != 'hunyuan_vl': + self.removeDupOps(self.mnn_model_path) self.mnn2json(self.mnn_model_path, mnn_json) if self.args.gptq_path is not None: self.apply_gptq(mnn_json) diff --git a/transformers/llm/export/utils/model.py b/transformers/llm/export/utils/model.py index e9030824e9..2e62fc305e 100644 --- a/transformers/llm/export/utils/model.py +++ b/transformers/llm/export/utils/model.py @@ -1,5 +1,7 @@ import torch import importlib +import json +import os from packaging.version import Version from transformers import PreTrainedModel, AutoConfig, AutoModel, AutoModelForCausalLM from typing import Optional, List @@ -9,6 +11,61 @@ from utils.model_mapper import ModelMapper from utils.transformers import Embedding, Rotary, Decoder, Lm + +def remap_hunyuan_vl_state_dict(state_dict): + remapped = {} + prefixes = ( + ('model.embed_tokens.', 'model.language_model.embed_tokens.'), + ('model.layers.', 'model.language_model.layers.'), + ('model.norm.', 'model.language_model.norm.'), + ('vit.embeddings.', 'model.vision_tower.embeddings.'), + ('vit.layers.', 'model.vision_tower.layers.'), + ('vit.perceive.before_rms.', 'model.vision_tower.patch_merger.before_rms.'), + ('vit.perceive.after_rms.', 'model.vision_tower.patch_merger.after_rms.'), + ('vit.perceive.image_begin', 'model.vision_tower.patch_merger.image_begin'), + ('vit.perceive.image_end', 'model.vision_tower.patch_merger.image_end'), + ('vit.perceive.image_newline', 'model.vision_tower.patch_merger.image_newline'), + ('vit.perceive.image_sep', 'model.vision_tower.patch_merger.image_sep'), + ('vit.perceive.mlp.', 'model.vision_tower.patch_merger.mlp.'), + ('vit.perceive.proj.0.', 'model.vision_tower.patch_merger.proj_conv.'), + ('vit.perceive.proj.2.', 'model.vision_tower.patch_merger.proj_out.'), + ) + for key, value in state_dict.items(): + new_key = key + for old_prefix, new_prefix in prefixes: + if key.startswith(old_prefix): + new_key = new_prefix + key[len(old_prefix):] + break + if new_key.startswith('model.vision_tower.layers.'): + new_key = new_key.replace('.input_layernorm.', '.layer_norm1.') + new_key = new_key.replace('.post_attention_layernorm.', '.layer_norm2.') + new_key = new_key.replace('.mlp.dense_h_to_4h.', '.mlp.fc1.') + new_key = new_key.replace('.mlp.dense_4h_to_h.', '.mlp.fc2.') + remapped[new_key] = value + return remapped + + +def load_hunyuan_vl_state_dict(model_path): + from safetensors.torch import load_file + index_path = os.path.join(model_path, 'model.safetensors.index.json') + if os.path.exists(index_path): + with open(index_path, 'r', encoding='utf-8') as f: + index = json.load(f) + shard_files = sorted(set(index.get('weight_map', {}).values())) + if not shard_files: + raise ValueError(f"HunyuanVL safetensors index has no weight_map entries: {index_path}") + else: + safetensors_path = os.path.join(model_path, 'model.safetensors') + if not os.path.exists(safetensors_path): + raise FileNotFoundError(f"HunyuanVL weights not found: {safetensors_path}") + shard_files = [os.path.basename(safetensors_path)] + state_dict = {} + for shard_file in shard_files: + shard_path = shard_file if os.path.isabs(shard_file) else os.path.join(model_path, shard_file) + state_dict.update(load_file(shard_path, device='cpu')) + return state_dict + + class LlmModel(PreTrainedModel): config_class = LlmConfig @@ -76,6 +133,7 @@ def get_model_class(model_type: str): 'glm_ocr': 'GlmOcrForConditionalGeneration', 'lfm2_vl': 'Lfm2VlForConditionalGeneration', 'gemma4': 'Gemma4ForConditionalGeneration', + 'hunyuan_vl': 'HunYuanVLForConditionalGeneration', } if model_type is None or model_type not in MODEL_CLASS_MAPPING: return AutoModelForCausalLM @@ -128,6 +186,22 @@ def from_pretrained(cls, pretrained_model_name_or_path, args=None, **kwargs): ) # Force sdpa attention on CPU (flash_attention_2 requires GPU) original_model.lfm.set_attn_implementation('sdpa') + elif model_type == 'hunyuan_vl': + try: + original_model = model_class.from_pretrained(pretrained_model_name_or_path, **load_kwargs) + except Exception as load_error: + original_config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True) + if hasattr(model_class, '_from_config'): + original_model = model_class._from_config(original_config) + else: + original_model = model_class(original_config) + state_dict = remap_hunyuan_vl_state_dict(load_hunyuan_vl_state_dict(pretrained_model_name_or_path)) + missing, unexpected = original_model.load_state_dict(state_dict, strict=False) + if missing or unexpected: + raise RuntimeError( + "HunyuanVL fallback weight load mismatch: " + f"missing={missing}, unexpected={unexpected}" + ) from load_error else: # Normal loading with weights try: @@ -493,7 +567,7 @@ def get_position_ids(self, seq_len, new_tokens=0, input_ids=None): position_ids = torch.arange(seq_len, dtype=torch.int) if self.rotary.is_mrope: - position_ids = torch.stack([position_ids] * 3) + position_ids = torch.stack([position_ids] * self.rotary.mrope_axes) else: position_ids = position_ids.unsqueeze(0) return position_ids diff --git a/transformers/llm/export/utils/model_mapper.py b/transformers/llm/export/utils/model_mapper.py index 30f7b88344..0fa51cc7a5 100644 --- a/transformers/llm/export/utils/model_mapper.py +++ b/transformers/llm/export/utils/model_mapper.py @@ -701,6 +701,40 @@ def regist_hunyuan_v1_dense(self): } self.regist('hunyuan_v1_dense', hunyuan_map) + def regist_hunyuan_vl(self): + hunyuan_vl_config = { + 'hidden_size': 'text_config.hidden_size', + 'head_dim': 'text_config.head_dim', + 'num_attention_heads': 'text_config.num_attention_heads', + 'num_hidden_layers': 'text_config.num_hidden_layers', + 'num_key_value_heads': 'text_config.num_key_value_heads', + 'rope_theta': 'text_config.rope_theta', + 'rope_scaling': 'text_config.rope_scaling', + 'max_position_embeddings': 'text_config.max_position_embeddings' + } + hunyuan_vl_model = { + 'lm': 'lm_head', + 'embed': 'model.language_model.embed_tokens', + 'blocks': 'model.language_model.layers', + 'final_layernorm': 'model.language_model.norm', + 'visual': 'model.vision_tower' + } + hunyuan_attention = { + 'q_proj': 'q_proj', + 'k_proj': 'k_proj', + 'v_proj': 'v_proj', + 'o_proj': 'o_proj', + 'q_norm': 'query_layernorm', + 'k_norm': 'key_layernorm' + } + hunyuan_vl_map = { + 'config': hunyuan_vl_config, + 'model': hunyuan_vl_model, + 'decoder': self.default_decoder, + 'attention': hunyuan_attention + } + self.regist('hunyuan_vl', hunyuan_vl_map) + def regist_gpt_oss(self): gpt_oss_config = { 'hidden_size': 'hidden_size', diff --git a/transformers/llm/export/utils/transformers.py b/transformers/llm/export/utils/transformers.py index 067b1b0e33..2c2d3b4527 100644 --- a/transformers/llm/export/utils/transformers.py +++ b/transformers/llm/export/utils/transformers.py @@ -947,6 +947,7 @@ def __init__(self, config): self.attention_scaling = 1.0 self.is_scaled = False self.mrope_interleaved = False + self.mrope_axes = 3 def get_theta(): return 1.0 / (self.rope_theta ** (torch.arange(0, self.rotary_dim, 2, dtype=torch.float32) / self.rotary_dim)) @@ -954,17 +955,21 @@ def get_theta(): self.theta = get_theta() # other type if hasattr(config, 'rope_scaling') and config.rope_scaling is not None: - scaling_config = config.rope_scaling + scaling_config = dict(config.rope_scaling) + if 'xdrope_section' in scaling_config and 'mrope_section' not in scaling_config: + scaling_config['mrope_section'] = scaling_config['xdrope_section'] # get rope_type rope_type = 'default' - if 'type' in config.rope_scaling: - rope_type = config.rope_scaling['type'] - elif 'rope_type' in config.rope_scaling: - rope_type = config.rope_scaling['rope_type'] + if 'type' in scaling_config: + rope_type = scaling_config['type'] + elif 'rope_type' in scaling_config: + rope_type = scaling_config['rope_type'] + if rope_type == 'xdrope': + rope_type = 'dynamic' # gen theta for rope_type if rope_type == 'dynamic': # NTK - if 'alpha' in config.rope_scaling: # NTKAlpha in Hunyuan - self.rope_theta *= (config.rope_scaling['alpha'] ** (self.rotary_dim / (self.rotary_dim - 2))) + if 'alpha' in scaling_config: # NTKAlpha in Hunyuan + self.rope_theta *= (scaling_config['alpha'] ** (self.rotary_dim / (self.rotary_dim - 2))) else: # NTKScaling pass self.theta = get_theta() @@ -989,6 +994,7 @@ def get_theta(): if 'mrope_section' in scaling_config: self.mrope_interleaved = scaling_config.get('mrope_interleaved', False) self.mrope_section = scaling_config['mrope_section'] + self.mrope_axes = len(self.mrope_section) self.theta = get_theta().unsqueeze(0) self.theta_sections = self.theta.split(self.mrope_section, dim=-1) def apply_interleaved_mrope(freqs, mrope_section): @@ -1028,17 +1034,26 @@ def mrope_forward(self, position_ids): idx_theta = position_ids * self.theta.to(position_ids.device) idx_theta = idx_theta.transpose(1, 0).reshape(-1, 3 * self.rotary_dim // 2) idx_theta = idx_theta[:, self.mrope_reindex] + elif self.model_type == 'hunyuan_vl': + axis_theta = position_ids * self.theta.to(position_ids.device) + axis_theta = torch.cat((axis_theta, axis_theta), dim=-1) + full_sections = [section * 2 for axis, section in enumerate(self.mrope_section)] + axis_chunks = [axis_theta[axis].split(full_sections, dim=-1) for axis in range(self.mrope_axes)] + idx_theta = torch.cat([ + axis_chunks[axis % self.mrope_axes][axis] for axis, section in enumerate(full_sections) + ], dim=-1) else: idx_theta = torch.concat([ - position_ids[0] * self.theta_sections[0], - position_ids[1] * self.theta_sections[1], - position_ids[2] * self.theta_sections[2] + position_ids[axis] * self.theta_sections[axis] + for axis, section in enumerate(self.mrope_section) ], dim=-1) rotary_pos_emb = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)]) if self.model_type in ['glm_ocr']: # interleaved doubling: [c0,c0,c1,c1,...,cn,cn] rotary_pos_emb = torch.stack((rotary_pos_emb, rotary_pos_emb), dim=-1) rotary_pos_emb = rotary_pos_emb.reshape(*rotary_pos_emb.shape[:-2], -1) + elif self.model_type == 'hunyuan_vl': + pass else: rotary_pos_emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) rotary_pos_emb = rotary_pos_emb.unsqueeze(2).unsqueeze(1) @@ -1490,4 +1505,4 @@ def forward(self, hidden_states): m_logits = m_logits / self.final_logit_softcapping m_logits = torch.tanh(m_logits) m_logits = m_logits * self.final_logit_softcapping - return m_logits \ No newline at end of file + return m_logits diff --git a/transformers/llm/export/utils/vision.py b/transformers/llm/export/utils/vision.py index ec975f2264..dbb1cd00f1 100644 --- a/transformers/llm/export/utils/vision.py +++ b/transformers/llm/export/utils/vision.py @@ -1,4 +1,6 @@ +import json import math +import os import torch import torch.nn.functional as F import numpy as np @@ -52,6 +54,7 @@ def get_vision(model_type): 'minicpmv': MiniCPMVision, 'glm_ocr': GlmOcrVision, 'lfm2_vl': Lfm2VlVision, + 'hunyuan_vl': HunyuanVLVision, } if model_type in visual_models: return visual_models[model_type] @@ -1467,6 +1470,271 @@ def export(self, onnx_path): }) return onnx_model + +class HunyuanVLVision(Vision): + def __init__(self, visual, base): + self.image_embeds = [] + self.image_grid_thw = [] + self.image_height = 512 + self.image_width = 512 + self.model_path = getattr(getattr(base, 'args', None), 'path', None) + self.is_mrope = getattr(getattr(base, 'rotary', None), 'is_mrope', False) + self.mrope_axes = getattr(getattr(base, 'rotary', None), 'mrope_axes', 3) + super().__init__(visual, base) + self.quant_bit = 4 + self.transformer_fuse = False + + def load(self): + vconfig = self.visual.config + self.vision_start_id = self.config.image_start_token_id + self.vision_end_id = self.config.image_end_token_id + self.image_pad_id = self.config.image_token_id + self.patch_size = vconfig.patch_size + self.merge_size = vconfig.spatial_merge_size + self.temporal_patch_size = vconfig.temporal_patch_size + self.min_pixels = vconfig.min_image_size * vconfig.min_image_size + self.max_pixels = vconfig.max_image_size * vconfig.max_image_size + if self.model_path is not None: + preprocessor_config = os.path.join(self.model_path, 'preprocessor_config.json') + if os.path.exists(preprocessor_config): + with open(preprocessor_config, 'r', encoding='utf-8') as f: + processor_config = json.load(f) + self.min_pixels = int(processor_config.get('min_pixels', self.min_pixels)) + self.max_pixels = int(processor_config.get('max_pixels', self.max_pixels)) + self.image_height = vconfig.min_image_size + self.image_width = vconfig.min_image_size + self.llm_config['vision_type'] = 'hunyuan_vl' + self.llm_config['image_size'] = self.image_height + self.llm_config['image_size_unit'] = self.patch_size * self.merge_size + self.llm_config['hunyuan_patch_size'] = self.patch_size + self.llm_config['hunyuan_spatial_merge_size'] = self.merge_size + self.llm_config['hunyuan_temporal_patch_size'] = self.temporal_patch_size + self.llm_config['image_max_size'] = vconfig.max_image_size + self.llm_config['image_min_pixels'] = self.min_pixels + self.llm_config['image_max_pixels'] = self.max_pixels + self.llm_config['vision_start'] = self.vision_start_id + self.llm_config['vision_end'] = self.vision_end_id + self.llm_config['image_pad'] = self.image_pad_id + self.vision_start_token = self.tokenizer.id_to_str(self.vision_start_id) + self.vision_end_token = self.tokenizer.id_to_str(self.vision_end_id) + self.image_pad_token = self.tokenizer.id_to_str(self.image_pad_id) + + def get_position_ids(self, input_ids, seq_len, new_tokens): + if not self.is_mrope: + return None + axes = self.mrope_axes + if new_tokens: + return torch.stack([torch.tensor([seq_len - 1], dtype=torch.int)] * axes) + position_ids = torch.arange(seq_len, dtype=torch.int) + position_ids = torch.stack([position_ids] * axes) + if input_ids is None or len(self.image_grid_thw) == 0: + return position_ids + image_token_id = getattr(self.config, 'image_token_id', None) + if image_token_id is None: + return position_ids + flat_input_ids = input_ids.reshape(-1).to(torch.int64) + image_mask = (flat_input_ids == int(image_token_id)).tolist() + spans = [] + start = None + for index, is_image in enumerate(image_mask): + if is_image and start is None: + start = index + elif not is_image and start is not None: + spans.append((start, index)) + start = None + if start is not None: + spans.append((start, len(image_mask))) + if len(spans) != len(self.image_grid_thw): + raise ValueError( + f"HunyuanVL image spans do not match image_grid_thw: spans={len(spans)}, " + f"grids={len(self.image_grid_thw)}" + ) + merge_size = int(getattr(self, 'merge_size', 1)) + axis_offset = max(0, axes - 3) + for image_index, ((span_start, span_end), grid_thw) in enumerate(zip(spans, self.image_grid_thw)): + _, grid_h, grid_w = [int(x) for x in grid_thw] + merged_h = grid_h // merge_size + merged_w = grid_w // merge_size + grid_tokens = merged_h * (merged_w + 1) + span_len = span_end - span_start + if span_len == grid_tokens + 2: + grid_start = span_start + 1 + elif span_len == grid_tokens: + grid_start = span_start + else: + raise ValueError( + "HunyuanVL image token span length does not match image_grid_thw: " + f"span_length={span_len}, expected {grid_tokens} or {grid_tokens + 2}" + ) + grid_end = grid_start + grid_tokens + width = torch.arange(merged_w + 1, dtype=torch.int).repeat(merged_h) + height = torch.arange(merged_h, dtype=torch.int).repeat_interleave(merged_w + 1) + position_ids[axis_offset, grid_start:grid_end] = width + position_ids[axis_offset + 1, grid_start:grid_end] = height + position_ids[axis_offset + 2, grid_start:grid_end] = image_index + return position_ids + + def str_to_ids(self, prompt): + self.image_embeds = [] + self.image_grid_thw = [] + if '' not in prompt or '' not in prompt: + return self.tokenizer(prompt, return_tensors="pt")['input_ids'] + import re + import requests + from PIL import Image + pattern = r'(.*?)' + parts = re.split(pattern, prompt) + txt_prompt = '' + for part in parts: + if re.match(pattern, part): + img_content = re.search(r'(.*?)', part).group(1) + image_hw = None + match = re.search(r'(.*?)', img_content) + if match: + img_content = img_content[:match.start()] + img_content[match.end():] + hw = match.group(1).split(',') + image_hw = (int(hw[0]), int(hw[1])) + if img_content.startswith('http://') or img_content.startswith('https://'): + image_obj = Image.open(requests.get(img_content, stream=True).raw) + else: + image_obj = Image.open(img_content) + img_pad_len = self.img_process(image_obj, image_hw) + txt_prompt += self.vision_start_token + txt_prompt += self.image_pad_token * img_pad_len + txt_prompt += self.vision_end_token + else: + txt_prompt += part + return self.tokenizer(txt_prompt, return_tensors="pt")['input_ids'] + + def smart_resize(self, height: int, width: int): + factor = self.patch_size * self.merge_size + if max(height, width) / min(height, width) > 200: + raise ValueError("absolute aspect ratio must be smaller than 200") + h_bar = round(height / factor) * factor + w_bar = round(width / factor) * factor + if h_bar * w_bar > self.max_pixels: + beta = math.sqrt((height * width) / self.max_pixels) + h_bar = max(factor, math.floor(height / beta / factor) * factor) + w_bar = max(factor, math.floor(width / beta / factor) * factor) + elif h_bar * w_bar < self.min_pixels: + beta = math.sqrt(self.min_pixels / (height * width)) + h_bar = math.ceil(height * beta / factor) * factor + w_bar = math.ceil(width * beta / factor) * factor + return h_bar, w_bar + + def vision_reshape(self, images): + batch, channel, height, width = images.shape + grid_h, grid_w = height // self.patch_size, width // self.patch_size + patches = images.reshape( + batch, + channel, + grid_h // self.merge_size, + self.merge_size, + self.patch_size, + grid_w // self.merge_size, + self.merge_size, + self.patch_size, + ) + patches = patches.permute(0, 2, 3, 5, 6, 1, 4, 7) + flatten_patches = patches.unsqueeze(6).expand( + -1, -1, -1, -1, -1, -1, self.temporal_patch_size, -1, -1 + ).reshape( + batch, + grid_h * grid_w, + channel * self.temporal_patch_size * self.patch_size * self.patch_size, + ) + grid_thw = torch.tensor([[1, grid_h, grid_w]], dtype=torch.long) + self.image_grid_thw.append([1, grid_h, grid_w]) + return flatten_patches.reshape(-1, flatten_patches.shape[-1]), grid_thw + + def images_forward(self, images): + pixel_values, image_grid_thw = self.vision_reshape(images) + return self.forward(pixel_values, image_grid_thw) + + def forward(self, pixel_values, image_grid_thw): + hidden_states = self.visual.embeddings(pixel_values, image_grid_thw) + for layer in self.visual.layers: + residual = hidden_states + hidden_states = layer.layer_norm1(hidden_states) + attn = layer.self_attn + batch_size, seq_len, _ = hidden_states.shape + query = attn.q_proj(hidden_states) + key = attn.k_proj(hidden_states) + value = attn.v_proj(hidden_states) + query = query.view(batch_size, seq_len, attn.num_heads, attn.head_dim).transpose(1, 2) + key = key.view(batch_size, seq_len, attn.num_heads, attn.head_dim).transpose(1, 2) + value = value.view(batch_size, seq_len, attn.num_heads, attn.head_dim).transpose(1, 2) + attn_weights = torch.matmul(query, key.transpose(2, 3)) * attn.scaling + attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + hidden_states = torch.matmul(attn_weights, value) + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, seq_len, -1).contiguous() + hidden_states = attn.o_proj(hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = layer.layer_norm2(hidden_states) + hidden_states = layer.mlp(hidden_states) + hidden_states = residual + hidden_states + output = self.visual.patch_merger(hidden_states, size=(image_grid_thw[0, 1], image_grid_thw[0, 2])) + return output.squeeze(0).unsqueeze(1) + + def img_process(self, image, image_hw=None): + from transformers.image_transforms import ( + convert_to_rgb, + resize, + rescale, + normalize + ) + from transformers.image_utils import ( + PILImageResampling, + infer_channel_dimension_format, + to_numpy_array + ) + image_width, image_height = image.size + if image_hw is not None: + image_height, image_width = image_hw + image = convert_to_rgb(image) + image = to_numpy_array(image) + resized_height, resized_width = self.smart_resize(image_height, image_width) + image_format = infer_channel_dimension_format(image) + image = resize( + image, + size=(resized_height, resized_width), + resample=PILImageResampling.BICUBIC, + input_data_format=image_format + ) + image = rescale(image, scale=1 / 255.0, input_data_format=image_format) + image = normalize(image=image, mean=self.norm_mean, std=self.norm_std, input_data_format=image_format) + image = np.expand_dims(image, [0]) + image = image.transpose(0, 3, 1, 2) + image_embed = self.images_forward(torch.from_numpy(image)) + self.image_embeds.append(image_embed.to(dtype=self.embed_.embed.weight.dtype)) + return image_embed.shape[0] + + def embed(self, input_ids, images=None, videos=None): + input_embeds = self.embed_(input_ids) + if self.image_embeds: + image_mask = (input_ids == self.image_pad_id).squeeze() + input_embeds[image_mask] = torch.concat(self.image_embeds, dim=0).to(input_embeds.dtype) + self.image_embeds = [] + return input_embeds + + @spinner_run(f'export visual to ') + def export(self, onnx_path): + grid = self.image_height // self.patch_size + pixel_values = torch.randn([grid * grid, 3 * self.temporal_patch_size * self.patch_size * self.patch_size]) + image_grid_thw = torch.tensor([[1, grid, grid]], dtype=torch.long) + onnx_model = f'{onnx_path}/visual.onnx' + onnx_export(self, (pixel_values, image_grid_thw), + onnx_model, + input_names=['pixel_values', 'image_grid_thw'], + output_names=['image_embeds'], + dynamic_axes={ + "pixel_values": { 0: "num_patches" }, + "image_grid_thw": { 0: "num_images" }, + }) + return onnx_model + # FastVLM class MobileCLIPVision(QwenVision): def __init__(self, visual, base):