2026-07-21 22:49:52 +08:00

46 lines
1.4 KiB
C++

#ifndef __SD_MODEL_IO_GGUF_IO_H__
#define __SD_MODEL_IO_GGUF_IO_H__
#include <string>
#include <vector>
#include "streaming_writer.h"
#include "tensor_storage.h"
struct ggml_context;
struct gguf_context;
bool is_gguf_file(const std::string& file_path);
bool read_gguf_file(const std::string& file_path,
std::vector<TensorStorage>& tensor_storages,
std::string* error = nullptr);
bool write_gguf_file(const std::string& file_path,
const std::vector<TensorWriteInfo>& tensors,
std::string* error = nullptr);
class GGUFStreamingWriter : public StreamingModelWriter {
public:
GGUFStreamingWriter() = default;
~GGUFStreamingWriter() override;
bool write_metadata(const std::string& file_path,
const std::vector<TensorWritePlan>& tensors,
std::string* error = nullptr) override;
bool write_tensor(std::ostream& output,
size_t tensor_index,
const uint8_t* data,
size_t size,
std::string* error = nullptr) const override;
uint64_t file_size() const override;
void close();
private:
std::vector<TensorWritePlan> tensors_;
std::vector<uint64_t> tensor_offsets_;
uint64_t file_size_ = 0;
ggml_context* meta_ctx_ = nullptr;
gguf_context* gguf_ctx_ = nullptr;
};
#endif // __SD_MODEL_IO_GGUF_IO_H__