Method loadModule

transformers/diffusion/engine/src/diffusion_sd35.cpp:114–160

}

bool DiffusionSD35::loadModule(int index) {
    if (mModules[index])
        return true;

    MNN::Express::Executor::getGlobalExecutor()->gc(MNN::Express::Executor::FULL);

    Module::Config module_config;
    module_config.shapeMutable = false;

    std::vector<std::string> modelNames = {"text_encoder.mnn", "text_encoder_2.mnn", "text_encoder_3.mnn",
                                           "transformer.mnn", "vae_decoder.mnn"};

    std::string model_path = mModelPath + "/" + modelNames[index];
    MNN_PRINT("Load %s\n", model_path.c_str());

    std::vector<std::string> inputs, outputs;
    if (index == 0) { // text_encoder
        inputs = {"input_ids"};
        outputs = {"last_hidden_state", "text_embeds"};
    } else if (index == 1) { // text_encoder_2
        inputs = {"input_ids"};
        outputs = {"last_hidden_state", "text_embeds"};
    } else if (index == 2) { // text_encoder_3
        inputs = {"input_ids"};
        outputs = {"last_hidden_state"};
    } else if (index == 3) { // transformer
        inputs = {"hidden_states", "encoder_hidden_states", "pooled_projections", "timestep"};
        outputs = {"out_hidden_states"};
    } else if (index == 4) { // vae_decoder
        inputs = {"latent_sample"};
        outputs = {"sample"};
    }

    mModules[index].reset(Module::load(inputs, outputs, model_path.c_str(), runtime_manager_, &module_config));
    if (!mModules[index]) {
        MNN_ERROR("Failed to load model %s\n", model_path.c_str());
        if (index == 2) {
            MNN_PRINT("Warning: T5 text encoder failed to load. Will use zeros.\n");
        } else {
            return false;
        }
    }

    if (mModules[index])
        mModules[index]->traceOrOptimize(MNN::Interpreter::Session_Resize_Fix);
    return true;
}

std::pair<VARP, VARP> DiffusionSD35::encode_prompt(const std::string& prompt) {
    auto run_encoder = [&](int module_index, VARP input_ids, const char* name) {