feat: stream model conversion (#1581)

src/model_io/gguf_io.cpp

@@ -1,7 +1,10 @@
 #include "gguf_io.h"
 
+#include <algorithm>
 #include <cstdint>
+#include <cstdio>
 #include <fstream>
+#include <ostream>
 #include <string>
 #include <vector>
 
@@ -121,3 +124,115 @@ bool write_gguf_file(const std::string& file_path,
     gguf_free(gguf_ctx);
     return success;
 }
+
+GGUFStreamingWriter::~GGUFStreamingWriter() {
+    close();
+}
+
+bool GGUFStreamingWriter::write_metadata(const std::string& file_path,
+                                         const std::vector<TensorWritePlan>& tensors,
+                                         std::string* error) {
+    close();
+    tensors_   = tensors;
+    file_size_ = 0;
+
+    size_t meta_mem = 1 * 1024 * 1024 + tensors.size() * ggml_tensor_overhead();
+    meta_ctx_       = ggml_init({meta_mem, nullptr, true});
+    if (meta_ctx_ == nullptr) {
+        set_error(error, "ggml_init failed for GGUF metadata");
+        return false;
+    }
+
+    gguf_ctx_ = gguf_init_empty();
+    if (gguf_ctx_ == nullptr) {
+        set_error(error, "gguf_init_empty failed");
+        close();
+        return false;
+    }
+
+    for (const TensorWritePlan& plan : tensors) {
+        ggml_tensor* tensor = ggml_new_tensor(meta_ctx_, plan.type, plan.n_dims, plan.ne);
+        if (tensor == nullptr) {
+            set_error(error, "ggml_new_tensor failed for tensor '" + plan.name + "'");
+            close();
+            return false;
+        }
+        ggml_set_name(tensor, plan.name.c_str());
+        gguf_add_tensor(gguf_ctx_, tensor);
+    }
+
+    LOG_INFO("trying to save tensors to %s", file_path.c_str());
+    FILE* file = fopen(file_path.c_str(), "wb+");
+    if (file == nullptr) {
+        set_error(error, "failed to open output file '" + file_path + "'");
+        close();
+        return false;
+    }
+
+    // ggml exposes GGUF metadata writing through FILE* only. Keep FILE usage
+    // isolated here; tensor data is written through std::fstream by the shared
+    // streaming pipeline.
+    if (!gguf_write_to_file_ptr(gguf_ctx_, file, true)) {
+        fclose(file);
+        set_error(error, "failed to write GGUF metadata to '" + file_path + "'");
+        close();
+        return false;
+    }
+    fclose(file);
+
+    const uint64_t data_start = gguf_get_meta_size(gguf_ctx_);
+    tensor_offsets_.resize(tensors.size());
+    file_size_ = data_start;
+    for (size_t i = 0; i < tensors.size(); i++) {
+        tensor_offsets_[i] = data_start + gguf_get_tensor_offset(gguf_ctx_, static_cast<int64_t>(i));
+        file_size_         = std::max(file_size_, tensor_offsets_[i] + tensors[i].nbytes());
+    }
+    return true;
+}
+
+bool GGUFStreamingWriter::write_tensor(std::ostream& output,
+                                       size_t tensor_index,
+                                       const uint8_t* data,
+                                       size_t size,
+                                       std::string* error) const {
+    if (tensor_index >= tensors_.size() || tensor_index >= tensor_offsets_.size()) {
+        set_error(error, "invalid GGUF tensor index");
+        return false;
+    }
+    const TensorWritePlan& plan = tensors_[tensor_index];
+    if (size != plan.nbytes()) {
+        set_error(error, "size mismatch while writing tensor '" + plan.name + "'");
+        return false;
+    }
+    output.seekp(static_cast<std::streamoff>(tensor_offsets_[tensor_index]), std::ios::beg);
+    if (!output) {
+        set_error(error, "failed to seek output for tensor '" + plan.name + "'");
+        return false;
+    }
+    if (size > 0) {
+        output.write(reinterpret_cast<const char*>(data), static_cast<std::streamsize>(size));
+    }
+    if (!output) {
+        set_error(error, "failed to write tensor '" + plan.name + "'");
+        return false;
+    }
+    return true;
+}
+
+uint64_t GGUFStreamingWriter::file_size() const {
+    return file_size_;
+}
+
+void GGUFStreamingWriter::close() {
+    tensor_offsets_.clear();
+    tensors_.clear();
+    file_size_ = 0;
+    if (gguf_ctx_ != nullptr) {
+        gguf_free(gguf_ctx_);
+        gguf_ctx_ = nullptr;
+    }
+    if (meta_ctx_ != nullptr) {
+        ggml_free(meta_ctx_);
+        meta_ctx_ = nullptr;
+    }
+}

src/model_io/gguf_io.h

@@ -4,8 +4,12 @@
 #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,
@@ -14,4 +18,28 @@ 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();
+
+    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__

src/model_io/safetensors_io.cpp

@@ -1,8 +1,10 @@
 #include "safetensors_io.h"
 
+#include <algorithm>
 #include <cstdint>
 #include <exception>
 #include <fstream>
+#include <ostream>
 #include <string>
 #include <vector>
 
@@ -314,3 +316,102 @@ bool write_safetensors_file(const std::string& file_path,
 
     return true;
 }
+
+bool SafetensorsStreamingWriter::write_metadata(const std::string& file_path,
+                                                const std::vector<TensorWritePlan>& tensors,
+                                                std::string* error) {
+    file_path_ = file_path;
+    tensors_   = tensors;
+    tensor_offsets_.clear();
+    data_start_ = 0;
+    file_size_  = 0;
+
+    nlohmann::ordered_json header = nlohmann::ordered_json::object();
+    uint64_t data_offset          = 0;
+    tensor_offsets_.resize(tensors.size());
+    for (size_t i = 0; i < tensors.size(); i++) {
+        const TensorWritePlan& plan = tensors[i];
+        std::string dtype;
+        if (!ggml_type_to_safetensors_dtype(plan.type, &dtype)) {
+            set_error(error,
+                      "unsupported safetensors dtype '" + std::string(ggml_type_name(plan.type)) +
+                          "' for tensor '" + plan.name + "'");
+            return false;
+        }
+
+        nlohmann::ordered_json json_tensor_info = nlohmann::ordered_json::object();
+        json_tensor_info["dtype"]               = dtype;
+
+        nlohmann::ordered_json shape = nlohmann::ordered_json::array();
+        for (int j = 0; j < plan.n_dims; ++j) {
+            shape.push_back(plan.ne[plan.n_dims - 1 - j]);
+        }
+        json_tensor_info["shape"] = shape;
+
+        nlohmann::ordered_json data_offsets = nlohmann::ordered_json::array();
+        data_offsets.push_back(data_offset);
+        data_offsets.push_back(data_offset + plan.nbytes());
+        json_tensor_info["data_offsets"] = data_offsets;
+
+        header[plan.name]  = json_tensor_info;
+        tensor_offsets_[i] = data_offset;
+        data_offset += plan.nbytes();
+    }
+
+    const std::string header_str = header.dump();
+    data_start_                  = ST_HEADER_SIZE_LEN + header_str.size();
+
+    LOG_INFO("trying to save tensors to %s", file_path.c_str());
+    std::ofstream file(file_path, std::ios::binary | std::ios::trunc);
+    if (!file.is_open()) {
+        set_error(error, "failed to open '" + file_path + "' for writing");
+        return false;
+    }
+
+    uint8_t header_size[ST_HEADER_SIZE_LEN];
+    for (int i = 0; i < static_cast<int>(ST_HEADER_SIZE_LEN); ++i) {
+        header_size[i] = static_cast<uint8_t>((header_str.size() >> (8 * i)) & 0xFF);
+    }
+    file.write(reinterpret_cast<const char*>(header_size), sizeof(header_size));
+    file.write(header_str.data(), static_cast<std::streamsize>(header_str.size()));
+    if (!file) {
+        set_error(error, "failed to write safetensors header to '" + file_path + "'");
+        return false;
+    }
+
+    file_size_ = data_start_ + data_offset;
+    return true;
+}
+
+bool SafetensorsStreamingWriter::write_tensor(std::ostream& output,
+                                              size_t tensor_index,
+                                              const uint8_t* data,
+                                              size_t size,
+                                              std::string* error) const {
+    if (tensor_index >= tensors_.size() || tensor_index >= tensor_offsets_.size()) {
+        set_error(error, "invalid safetensors tensor index");
+        return false;
+    }
+    const TensorWritePlan& plan = tensors_[tensor_index];
+    if (size != plan.nbytes()) {
+        set_error(error, "size mismatch while writing tensor '" + plan.name + "'");
+        return false;
+    }
+    output.seekp(static_cast<std::streamoff>(data_start_ + tensor_offsets_[tensor_index]), std::ios::beg);
+    if (!output) {
+        set_error(error, "failed to seek output for tensor '" + plan.name + "'");
+        return false;
+    }
+    if (size > 0) {
+        output.write(reinterpret_cast<const char*>(data), static_cast<std::streamsize>(size));
+    }
+    if (!output) {
+        set_error(error, "failed to write tensor '" + plan.name + "' to '" + file_path_ + "'");
+        return false;
+    }
+    return true;
+}
+
+uint64_t SafetensorsStreamingWriter::file_size() const {
+    return file_size_;
+}

src/model_io/safetensors_io.h

@@ -4,6 +4,7 @@
 #include <string>
 #include <vector>
 
+#include "streaming_writer.h"
 #include "tensor_storage.h"
 
 bool is_safetensors_file(const std::string& file_path);
@@ -14,4 +15,26 @@ bool write_safetensors_file(const std::string& file_path,
                             const std::vector<TensorWriteInfo>& tensors,
                             std::string* error = nullptr);
 
+class SafetensorsStreamingWriter : public StreamingModelWriter {
+public:
+    SafetensorsStreamingWriter() = default;
+
+    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;
+
+private:
+    std::string file_path_;
+    std::vector<TensorWritePlan> tensors_;
+    std::vector<uint64_t> tensor_offsets_;
+    uint64_t data_start_ = 0;
+    uint64_t file_size_  = 0;
+};
+
 #endif  // __SD_MODEL_IO_SAFETENSORS_IO_H__

src/model_io/streaming_writer.h

@@ -0,0 +1,26 @@
+#ifndef __SD_MODEL_IO_STREAMING_WRITER_H__
+#define __SD_MODEL_IO_STREAMING_WRITER_H__
+
+#include <cstdint>
+#include <iosfwd>
+#include <string>
+#include <vector>
+
+#include "tensor_storage.h"
+
+class StreamingModelWriter {
+public:
+    virtual ~StreamingModelWriter() = default;
+
+    virtual bool write_metadata(const std::string& file_path,
+                                const std::vector<TensorWritePlan>& tensors,
+                                std::string* error = nullptr)     = 0;
+    virtual bool write_tensor(std::ostream& output,
+                              size_t tensor_index,
+                              const uint8_t* data,
+                              size_t size,
+                              std::string* error = nullptr) const = 0;
+    virtual uint64_t file_size() const                            = 0;
+};
+
+#endif  // __SD_MODEL_IO_STREAMING_WRITER_H__

src/model_io/tensor_storage.h

@@ -127,6 +127,25 @@ struct TensorWriteInfo {
     ggml_tensor* tensor     = nullptr;
 };
 
+struct TensorWritePlan {
+    std::string name;
+    ggml_type type          = GGML_TYPE_F32;
+    int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
+    int n_dims              = 0;
+
+    int64_t nelements() const {
+        int64_t n = 1;
+        for (int i = 0; i < SD_MAX_DIMS; i++) {
+            n *= ne[i];
+        }
+        return n;
+    }
+
+    uint64_t nbytes() const {
+        return nelements() * ggml_type_size(type) / ggml_blck_size(type);
+    }
+};
+
 typedef std::function<bool(const TensorStorage&, ggml_tensor**)> on_new_tensor_cb_t;
 
 #endif  // __SD_TENSOR_STORAGE_H__