Method encodePrompt

transformers/diffusion/engine/src/wan_diffusion.cpp:122–175

}

VARP WanDiffusion::encodePrompt(const std::string& prompt, int* seqLen, VARP* outMask) {
    if (!mTokenizer || mModules.size() < 1 || !mModules[0]) {
        MNN_ERROR("Wan text encoder is not ready\n");
        return nullptr;
    }

    // Defensive: peek at the un-padded token count first so we can warn the
    // caller when the prompt is truncated. mTokenizer->encode pads/truncates
    // silently when maxlen > 0; passing maxlen=0 returns the raw ids.
    std::vector<int> rawTokens = mTokenizer->encode(prompt, 0);
    if ((int)rawTokens.size() > mMaxTextLen) {
        MNN_PRINT("Wan: prompt truncated from %d to %d tokens\n", (int)rawTokens.size(), mMaxTextLen);
    }

    std::vector<int> uncond = mTokenizer->encode("", mMaxTextLen);
    std::vector<int> cond = mTokenizer->encode(prompt, mMaxTextLen);
    if ((int)uncond.size() < mMaxTextLen || (int)cond.size() < mMaxTextLen) {
        MNN_ERROR("Wan tokenizer returned too few ids: uncond=%d cond=%d expected=%d\n", (int)uncond.size(),
                  (int)cond.size(), mMaxTextLen);
        return nullptr;
    }

    VARP inputIds = _Input({2, mMaxTextLen}, NCHW, halide_type_of<int>());
    int* inputPtr = inputIds->writeMap<int>();
    ::memset(inputPtr, 0, 2 * mMaxTextLen * sizeof(int));
    ::memcpy(inputPtr, uncond.data(), mMaxTextLen * sizeof(int));
    ::memcpy(inputPtr + mMaxTextLen, cond.data(), mMaxTextLen * sizeof(int));

    // Build attention mask from token ids: 1 where id != 0, 0 for padding.
    // Dtype must be int32 to match the ONNX export (see wan_onnx_export.py).
    VARP mask = _Input({2, mMaxTextLen}, NCHW, halide_type_of<int>());
    int* maskData = mask->writeMap<int>();
    for (int i = 0; i < 2 * mMaxTextLen; ++i) {
        maskData[i] = (inputPtr[i] != 0) ? 1 : 0;
    }
    if (outMask) {
        *outMask = mask;
    }

    auto outputs = mModules[0]->onForward({inputIds});
    if (outputs.empty() || outputs[0].get() == nullptr) {
        MNN_ERROR("Wan text_encoder returned empty output\n");
        return nullptr;
    }
    auto hiddenStates = _Convert(outputs[0], NCHW);
    auto info = hiddenStates->getInfo();
    if (seqLen != nullptr && info != nullptr && info->dim.size() > 1) {
        *seqLen = info->dim[1];
    } else if (seqLen != nullptr) {
        *seqLen = mMaxTextLen;
    }
    hiddenStates.fix(VARP::CONSTANT);
    return hiddenStates;
}

VARP WanDiffusion::transformer(VARP hiddenStates, VARP timestep, VARP encoderHiddenStates, VARP encoderAttentionMask) {
    auto outputs = mModules[1]->onForward({hiddenStates, timestep, encoderHiddenStates, encoderAttentionMask});