Method saveVideoFrames

transformers/diffusion/engine/src/wan_diffusion.cpp:212–272

}

bool WanDiffusion::saveVideoFrames(VARP sample, const std::string& outputDir, int requestedFrames) {
    if (sample.get() == nullptr) {
        MNN_ERROR("Wan VAE output is null\n");
        return false;
    }
    if (!MNNCreateDir(outputDir.c_str())) {
        MNN_ERROR("Failed to create output dir %s\n", outputDir.c_str());
        return false;
    }

    auto info = sample->getInfo();
    if (info == nullptr) {
        MNN_ERROR("Wan VAE output has no shape info\n");
        return false;
    }

    int saved = 0;
    if (info->dim.size() == 4) {
        MNN_PRINT("Warning: Wan VAE returned 4D/single-frame output; saving only frame_0000.png\n");
        std::string fileName = MNNFilePathConcat(outputDir, "frame_0000.png");
        return saveFrame(sample, fileName);
    }

    if (info->dim.size() != 5) {
        MNN_ERROR("Wan VAE output rank %d is not supported for frame saving\n", (int)info->dim.size());
        return false;
    }

    if (info->dim[1] != 3) {
        MNN_PRINT(
            "Warning: Wan VAE 5D output shape is not NCTHW with C=3; frame saving is not implemented for "
            "this layout\n");
        return false;
    }

    int decodedFrames = info->dim[2];
    int saveFrames = std::min(decodedFrames, requestedFrames);
    if (decodedFrames < requestedFrames) {
        MNN_PRINT("Warning: Wan VAE returned %d frames, requested %d; saving decoded frames only\n", decodedFrames,
                  requestedFrames);
    }
    for (int i = 0; i < saveFrames; ++i) {
        int startsData[5] = {0, 0, i, 0, 0};
        int sizesData[5] = {1, 3, 1, info->dim[3], info->dim[4]};
        auto frame = _Slice(sample, _Const(startsData, {5}, NCHW, halide_type_of<int>()),
                            _Const(sizesData, {5}, NCHW, halide_type_of<int>()));
        frame = _Reshape(frame, {1, 3, info->dim[3], info->dim[4]}, NCHW);

        std::ostringstream name;
        name << "frame_" << std::setfill('0') << std::setw(4) << i << ".png";
        std::string fileName = MNNFilePathConcat(outputDir, name.str());
        if (saveFrame(frame, fileName)) {
            ++saved;
        } else {
            MNN_PRINT("Warning: failed to save %s\n", fileName.c_str());
        }
    }