stable-diffusion.cpp/qcom_ml_vae_bridge.hpp at work/main · happyyzy/stable-diffusion.cpp · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
#ifndef SD_QCOM_ML_VAE_BRIDGE_HPP
#define SD_QCOM_ML_VAE_BRIDGE_HPP

#include <string>
#include <vector>

struct QcomMlVaeConfig {
    std::string model_dir;
    bool disable_mnn_attention   = true;
    bool fallback_attn_len_16384 = true;
};

class QcomMlVaeBridge {
public:
    QcomMlVaeBridge() = default;
    ~QcomMlVaeBridge();

    bool init(const std::string& shared_lib_path, const QcomMlVaeConfig& cfg, std::string* err_msg);

    bool decode(const std::vector<float>& latent_nchw,
                int latent_w,
                int latent_h,
                int latent_c,
                int batch,
                std::vector<float>* out_nchw,
                int out_w,
                int out_h,
                int out_c,
                std::string* err_msg,
                bool* used_attn_fallback_16384);

    bool prepare(int latent_w,
                 int latent_h,
                 int latent_c,
                 int batch,
                 std::string* err_msg);

    bool ready() const;

private:
    QcomMlVaeBridge(const QcomMlVaeBridge&)            = delete;
    QcomMlVaeBridge& operator=(const QcomMlVaeBridge&) = delete;

    void* m_lib_handle = nullptr;
    void* m_ctx_handle = nullptr;

    typedef int (*CreateFn)(const char* model_dir,
                            int disable_mnn_attention,
                            int fallback_attn_len_16384,
                            void** out_ctx,
                            char* err_buf,
                            int err_buf_cap);
    typedef int (*DecodeFn)(void* ctx,
                            const float* latent_nchw,
                            int latent_w,
                            int latent_h,
                            int latent_c,
                            int batch,
                            float* out_nchw,
                            int out_w,
                            int out_h,
                            int out_c,
                            int* used_attn_fallback_16384,
                            char* err_buf,
                            int err_buf_cap);
    typedef int (*PrepareFn)(void* ctx,
                             int latent_w,
                             int latent_h,
                             int latent_c,
                             int batch,
                             char* err_buf,
                             int err_buf_cap);
    typedef void (*DestroyFn)(void* ctx);

    CreateFn  m_create_fn  = nullptr;
    DecodeFn  m_decode_fn  = nullptr;
    PrepareFn m_prepare_fn = nullptr;
    DestroyFn m_destroy_fn = nullptr;
};

#endif  // SD_QCOM_ML_VAE_BRIDGE_HPP