Method unet

transformers/diffusion/engine/src/stable_diffusion.cpp:235–298

}

VARP StableDiffusion::unet(VARP text_embeddings, int iterNum, int randomSeed, std::function<void(int)> progressCallback) {
    if(mMemoryMode != 1) {
        mModules[0].reset();
    }
    if(mInitNoise.size() != 16384) {
        mInitNoise.resize(16384);
    }
#ifdef MNN_DUMP_DATA
    std::ostringstream fileName;
    fileName << "random.txt";
    std::ifstream input(fileName.str().c_str());
    for (int i = 0; i < 16384; ++i) {
        input >> mInitNoise[i];
    }
#else
    int seed = randomSeed < 0 ? std::random_device()() : randomSeed;
    std::mt19937 rng;
    rng.seed(seed);
    
    std::normal_distribution<float> normal(0, 1);
    for (int i = 0; i < 16384; i++) {
        mInitNoise[i] = normal(rng);
    }
#endif

    memcpy((void *)mLatentVar->writeMap<int8_t>(), mInitNoise.data(), 16384*sizeof(float));
    
    VARP scalevar = _Input({1}, NCHW, halide_type_of<float>());
    auto scaleptr = scalevar->writeMap<float>();
    scaleptr[0] = 7.5;

    auto floatVar = _Input({1}, NCHW, halide_type_of<float>());
    auto ptr = floatVar->writeMap<float>();
    auto plms = mLatentVar;

    for (int i = 0; i < mTimeSteps.size(); i++) {
        AUTOTIME;

        int timestep = mTimeSteps[i];
        ptr[0] = timestep;
        auto temp = _Cast(floatVar, halide_type_of<int>());
        mTimestepVar->input(temp);

        mSampleVar = _Concat({plms, plms}, 0);
        auto outputs = forwardWithResizeCache(1, {mSampleVar, mTimestepVar, text_embeddings});
        auto output = _Convert(outputs[0], NCHW);

        auto noise_pred = output;

        auto splitvar = _Split(noise_pred, {2}, 0);
        auto noise_pred_uncond = splitvar[0];
        auto noise_pred_text = splitvar[1];

        noise_pred = scalevar * (noise_pred_text - noise_pred_uncond) + noise_pred_uncond;

        plms = step_plms(plms, noise_pred, i);

        if (progressCallback) {