feat: stream model conversion (#1581)

examples/cli/main.cpp

@@ -653,7 +653,8 @@ int main(int argc, const char* argv[]) {
                                                cli_params.output_path.c_str(),
                                                ctx_params.wtype,
                                                ctx_params.tensor_type_rules.c_str(),
-                                               cli_params.convert_name);
+                                               cli_params.convert_name,
+                                               ctx_params.n_threads);
         if (!success) {
             LOG_ERROR("convert '%s'/'%s' to '%s' failed",
                       ctx_params.model_path.c_str(),

examples/server/routes_sdapi.cpp

@@ -281,7 +281,7 @@ static nlohmann::json prepare_info_field(const SDContextParams& ctx_params,
         jsoninfo["clip_skip"] = gen_params.clip_skip;
     }
     if (gen_params.sample_params.scheduler != scheduler_t::SCHEDULER_COUNT) {
-        jsoninfo["extra_generation_params"] = nlohmann::json::object();
+        jsoninfo["extra_generation_params"]                  = nlohmann::json::object();
         jsoninfo["extra_generation_params"]["Schedule type"] = sd_scheduler_name(gen_params.sample_params.scheduler);
     }
     if (img2img) {
@@ -363,18 +363,18 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
                     continue;
                 }
 
-                bool embed_meta    = request.gen_params.embed_image_metadata;
+                bool embed_meta = request.gen_params.embed_image_metadata;
 
                 std::string params = get_image_params(*runtime->ctx_params,
                                                       request.gen_params,
                                                       request.gen_params.seed + i / images_per_batch);
 
-                auto image_bytes   = encode_image_to_vector(EncodedImageFormat::PNG,
-                                                            results[i].data,
-                                                            results[i].width,
-                                                            results[i].height,
-                                                            results[i].channel,
-                                                            embed_meta ? params : "");
+                auto image_bytes = encode_image_to_vector(EncodedImageFormat::PNG,
+                                                          results[i].data,
+                                                          results[i].width,
+                                                          results[i].height,
+                                                          results[i].channel,
+                                                          embed_meta ? params : "");
 
                 if (image_bytes.empty()) {
                     LOG_ERROR("write image to mem failed");

include/stable-diffusion.h

@@ -521,7 +521,8 @@ SD_API bool convert_with_components(const char* model_path,
                                     const char* output_path,
                                     enum sd_type_t output_type,
                                     const char* tensor_type_rules,
-                                    bool convert_name);
+                                    bool convert_name,
+                                    int n_threads);
 
 SD_API bool preprocess_canny(sd_image_t image,
                              float high_threshold,

src/convert.cpp

@@ -1,14 +1,33 @@
+#include <algorithm>
+#include <condition_variable>
+#include <cstdint>
 #include <cstring>
+#include <exception>
+#include <fstream>
+#include <memory>
 #include <mutex>
 #include <regex>
+#include <string>
+#include <thread>
 #include <vector>
 
+#include "core/util.h"
 #include "model_io/gguf_io.h"
 #include "model_io/safetensors_io.h"
+#include "model_io/streaming_writer.h"
 #include "model_loader.h"
-#include "util.h"
 
-#include "ggml_extend_backend.h"
+struct TensorExportInfo {
+    TensorStorage storage;
+    ggml_type type;
+};
+
+struct TensorExportJob {
+    TensorExportInfo info;
+    std::vector<uint8_t> data;
+    std::string error;
+    bool success = false;
+};
 
 static ggml_type get_export_tensor_type(ModelLoader& model_loader,
                                         const TensorStorage& tensor_storage,
@@ -33,47 +52,262 @@ static ggml_type get_export_tensor_type(ModelLoader& model_loader,
     return tensor_type;
 }
 
-static bool load_tensors_for_export(ModelLoader& model_loader,
-                                    ggml_context* ggml_ctx,
-                                    ggml_type type,
-                                    const TensorTypeRules& tensor_type_rules,
-                                    std::vector<TensorWriteInfo>& tensors) {
-    std::mutex tensor_mutex;
-    auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
-        const std::string& name = tensor_storage.name;
-        ggml_type tensor_type   = get_export_tensor_type(model_loader, tensor_storage, type, tensor_type_rules);
+static bool collect_tensors_for_export(ModelLoader& model_loader,
+                                       ggml_type type,
+                                       const TensorTypeRules& tensor_type_rules,
+                                       std::vector<TensorExportInfo>& tensors) {
+    tensors.clear();
+    tensors.reserve(model_loader.get_tensor_storage_map().size());
+    for (const auto& kv : model_loader.get_tensor_storage_map()) {
+        const TensorStorage& tensor_storage = kv.second;
+        TensorExportInfo info;
+        info.storage = tensor_storage;
+        info.type    = get_export_tensor_type(model_loader, tensor_storage, type, tensor_type_rules);
+        tensors.push_back(std::move(info));
+    }
+    LOG_INFO("collected %zu tensors for export", tensors.size());
+    return true;
+}
 
-        std::lock_guard<std::mutex> lock(tensor_mutex);
-        ggml_tensor* tensor = ggml_new_tensor(ggml_ctx, tensor_type, tensor_storage.n_dims, tensor_storage.ne);
-        if (tensor == nullptr) {
-            LOG_ERROR("ggml_new_tensor failed");
+static size_t export_tensor_nbytes(const TensorExportInfo& info) {
+    TensorStorage output_storage = info.storage;
+    output_storage.type          = info.type;
+    return static_cast<size_t>(output_storage.nbytes());
+}
+
+static TensorWritePlan tensor_write_plan_from_export_info(const TensorExportInfo& info) {
+    TensorWritePlan plan;
+    plan.name   = info.storage.name;
+    plan.type   = info.type;
+    plan.n_dims = info.storage.n_dims;
+    for (int i = 0; i < SD_MAX_DIMS; i++) {
+        plan.ne[i] = info.storage.ne[i];
+    }
+    return plan;
+}
+
+static std::vector<TensorWritePlan> tensor_write_plans_from_export_infos(const std::vector<TensorExportInfo>& tensors) {
+    std::vector<TensorWritePlan> plans;
+    plans.reserve(tensors.size());
+    for (const TensorExportInfo& info : tensors) {
+        plans.push_back(tensor_write_plan_from_export_info(info));
+    }
+    return plans;
+}
+
+static bool preallocate_output_file(const std::string& output_path, uint64_t file_size, std::string* error) {
+    if (file_size == 0) {
+        return true;
+    }
+
+    std::fstream file(output_path, std::ios::binary | std::ios::in | std::ios::out);
+    if (!file.is_open()) {
+        if (error != nullptr) {
+            *error = "failed to open output file '" + output_path + "' for preallocation";
+        }
+        return false;
+    }
+
+    // This portable fallback sets the final file size. A platform-specific
+    // posix_fallocate/ftruncate path can replace it later.
+    file.seekp(static_cast<std::streamoff>(file_size - 1), std::ios::beg);
+    file.put('\0');
+    file.flush();
+    if (!file) {
+        if (error != nullptr) {
+            *error = "failed to preallocate output file '" + output_path + "'";
+        }
+        return false;
+    }
+    return true;
+}
+
+static bool load_tensor_for_export(ModelLoader& model_loader, TensorExportJob& job) {
+    size_t mem_size = 1 * 1024 * 1024;
+    mem_size += ggml_tensor_overhead();
+    TensorStorage output_storage = job.info.storage;
+    output_storage.type          = job.info.type;
+    mem_size += static_cast<size_t>(output_storage.nbytes());
+
+    ggml_context* ggml_ctx = ggml_init({mem_size, nullptr, false});
+    if (ggml_ctx == nullptr) {
+        job.error = "ggml_init failed for tensor '" + job.info.storage.name + "'";
+        return false;
+    }
+
+    ggml_tensor* tensor = ggml_new_tensor(ggml_ctx, job.info.type, job.info.storage.n_dims, job.info.storage.ne);
+    if (tensor == nullptr) {
+        ggml_free(ggml_ctx);
+        job.error = "ggml_new_tensor failed for tensor '" + job.info.storage.name + "'";
+        return false;
+    }
+    ggml_set_name(tensor, job.info.storage.name.c_str());
+
+    const size_t tensor_nbytes = ggml_nbytes(tensor);
+    if (tensor_nbytes > 0 && !model_loader.load_tensor(job.info.storage, tensor)) {
+        ggml_free(ggml_ctx);
+        job.error = "failed to load tensor '" + job.info.storage.name + "'";
+        return false;
+    }
+
+    job.data.resize(tensor_nbytes);
+    if (tensor_nbytes > 0) {
+        memcpy(job.data.data(), tensor->data, tensor_nbytes);
+    }
+    ggml_free(ggml_ctx);
+    return true;
+}
+
+static bool stream_tensor_data(ModelLoader& model_loader,
+                               const std::string& output_path,
+                               const std::vector<TensorExportInfo>& tensors,
+                               const StreamingModelWriter& writer,
+                               int n_threads,
+                               std::string* error) {
+    n_threads = n_threads > 0 ? n_threads : sd_get_num_physical_cores();
+    n_threads = std::max(1, n_threads);
+    LOG_INFO("streaming convert with %d threads", n_threads);
+
+    int64_t start_time       = ggml_time_ms();
+    uint64_t bytes_written   = 0;
+    size_t tensors_written   = 0;
+    size_t next_tensor_index = 0;
+    bool failed              = false;
+    std::string failure;
+
+    const size_t memory_budget = 1024ull * 1024ull * 1024ull;
+    size_t reserved_bytes      = 0;
+
+    std::mutex work_mutex;
+    std::mutex progress_mutex;
+    std::condition_variable memory_cv;
+    std::vector<std::thread> workers;
+    workers.reserve(n_threads);
+
+    auto reserve_memory = [&](size_t bytes) -> bool {
+        std::unique_lock<std::mutex> lock(work_mutex);
+        memory_cv.wait(lock, [&]() {
+            return failed || reserved_bytes == 0 || reserved_bytes + bytes <= memory_budget;
+        });
+        if (failed) {
             return false;
         }
-        ggml_set_name(tensor, name.c_str());
-
-        if (!tensor->data) {
-            GGML_ASSERT(ggml_nelements(tensor) == 0);
-            // Avoid crashing writers by setting a dummy pointer for zero-sized tensors.
-            LOG_DEBUG("setting dummy pointer for zero-sized tensor %s", name.c_str());
-            tensor->data = ggml_get_mem_buffer(ggml_ctx);
-        }
-
-        TensorWriteInfo write_info;
-        write_info.tensor = tensor;
-        write_info.n_dims = tensor_storage.n_dims;
-        for (int i = 0; i < tensor_storage.n_dims; ++i) {
-            write_info.ne[i] = tensor_storage.ne[i];
-        }
-
-        *dst_tensor = tensor;
-        tensors.push_back(std::move(write_info));
-
+        reserved_bytes += bytes;
         return true;
     };
 
-    bool success = model_loader.load_tensors(on_new_tensor_cb);
-    LOG_INFO("load tensors done");
-    return success;
+    auto release_memory = [&](size_t bytes) {
+        {
+            std::lock_guard<std::mutex> lock(work_mutex);
+            reserved_bytes -= std::min(reserved_bytes, bytes);
+        }
+        memory_cv.notify_all();
+    };
+
+    auto fail = [&](const std::string& message) {
+        {
+            std::lock_guard<std::mutex> lock(work_mutex);
+            if (!failed) {
+                failed  = true;
+                failure = message;
+            }
+        }
+        memory_cv.notify_all();
+    };
+
+    for (int worker = 0; worker < n_threads; worker++) {
+        workers.emplace_back([&]() {
+            std::fstream output_file(output_path, std::ios::binary | std::ios::in | std::ios::out);
+            if (!output_file.is_open()) {
+                fail("failed to open output file '" + output_path + "' for tensor writing");
+                return;
+            }
+
+            while (true) {
+                size_t tensor_index = 0;
+                {
+                    std::lock_guard<std::mutex> lock(work_mutex);
+                    if (failed || next_tensor_index >= tensors.size()) {
+                        return;
+                    }
+                    tensor_index = next_tensor_index++;
+                }
+
+                const size_t tensor_bytes = export_tensor_nbytes(tensors[tensor_index]);
+                if (!reserve_memory(tensor_bytes)) {
+                    return;
+                }
+
+                TensorExportJob job;
+                job.info = tensors[tensor_index];
+                try {
+                    job.success = load_tensor_for_export(model_loader, job);
+                } catch (const std::exception& e) {
+                    job.error   = e.what();
+                    job.success = false;
+                }
+
+                if (!job.success) {
+                    release_memory(tensor_bytes);
+                    fail(job.error.empty() ? "streaming conversion failed" : job.error);
+                    return;
+                }
+
+                std::string write_error;
+                if (!writer.write_tensor(output_file,
+                                         tensor_index,
+                                         job.data.empty() ? nullptr : job.data.data(),
+                                         job.data.size(),
+                                         &write_error)) {
+                    release_memory(tensor_bytes);
+                    fail(write_error.empty() ? "streaming conversion write failed" : write_error);
+                    return;
+                }
+
+                {
+                    std::lock_guard<std::mutex> lock(progress_mutex);
+                    bytes_written += job.data.size();
+                    tensors_written++;
+                    float elapsed_seconds = (ggml_time_ms() - start_time) / 1000.0f;
+                    pretty_bytes_progress(static_cast<int>(tensors_written),
+                                          static_cast<int>(tensors.size()),
+                                          bytes_written,
+                                          elapsed_seconds);
+                }
+                release_memory(tensor_bytes);
+            }
+        });
+    }
+
+    for (auto& worker : workers) {
+        worker.join();
+    }
+    printf("\n");
+    if (failed) {
+        if (error != nullptr) {
+            *error = failure;
+        }
+        return false;
+    }
+    LOG_INFO("streaming conversion completed, taking %.2fs", (ggml_time_ms() - start_time) / 1000.f);
+    return true;
+}
+
+static bool write_model_file_streaming(ModelLoader& model_loader,
+                                       const std::string& output_path,
+                                       const std::vector<TensorExportInfo>& tensors,
+                                       StreamingModelWriter& writer,
+                                       int n_threads,
+                                       std::string* error) {
+    std::vector<TensorWritePlan> plans = tensor_write_plans_from_export_infos(tensors);
+    if (!writer.write_metadata(output_path, plans, error)) {
+        return false;
+    }
+    if (!preallocate_output_file(output_path, writer.file_size(), error)) {
+        return false;
+    }
+    model_loader.process_model_files(false, false);
+    return stream_tensor_data(model_loader, output_path, tensors, writer, n_threads, error);
 }
 
 static bool init_convert_path(ModelLoader& model_loader, const char* path, const char* prefix, bool& loaded_any) {
@@ -91,42 +325,29 @@ static bool init_convert_path(ModelLoader& model_loader, const char* path, const
 static bool export_loaded_model(ModelLoader& model_loader,
                                 const char* output_path,
                                 sd_type_t output_type,
-                                const char* tensor_type_rules) {
+                                const char* tensor_type_rules,
+                                int n_threads) {
     ggml_type type             = sd_type_to_ggml_type(output_type);
     bool output_is_safetensors = ends_with(output_path, ".safetensors");
     TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules);
 
-    auto backend    = sd_backend_cpu_init();
-    size_t mem_size = 1 * 1024 * 1024;  // for padding
-    mem_size += model_loader.get_tensor_storage_map().size() * ggml_tensor_overhead();
-    mem_size += model_loader.get_params_mem_size(backend, type);
-    LOG_INFO("model tensors mem size: %.2fMB", mem_size / 1024.f / 1024.f);
-    ggml_context* ggml_ctx = ggml_init({mem_size, nullptr, false});
-
-    if (ggml_ctx == nullptr) {
-        LOG_ERROR("ggml_init failed for converter");
-        ggml_backend_free(backend);
-        return false;
-    }
-
-    std::vector<TensorWriteInfo> tensors;
-    bool success = load_tensors_for_export(model_loader, ggml_ctx, type, type_rules, tensors);
-    ggml_backend_free(backend);
-
+    std::vector<TensorExportInfo> tensors;
+    bool success = collect_tensors_for_export(model_loader, type, type_rules, tensors);
     std::string error;
     if (success) {
+        std::unique_ptr<StreamingModelWriter> writer;
         if (output_is_safetensors) {
-            success = write_safetensors_file(output_path, tensors, &error);
+            writer = std::make_unique<SafetensorsStreamingWriter>();
         } else {
-            success = write_gguf_file(output_path, tensors, &error);
+            writer = std::make_unique<GGUFStreamingWriter>();
         }
+        success = write_model_file_streaming(model_loader, output_path, tensors, *writer, n_threads, &error);
     }
 
     if (!success && !error.empty()) {
         LOG_ERROR("%s", error.c_str());
     }
 
-    ggml_free(ggml_ctx);
     return success;
 }
 
@@ -139,7 +360,8 @@ bool convert_with_components(const char* model_path,
                              const char* output_path,
                              sd_type_t output_type,
                              const char* tensor_type_rules,
-                             bool convert_name) {
+                             bool convert_name,
+                             int n_threads) {
     ModelLoader model_loader;
     bool loaded_any = false;
 
@@ -161,7 +383,7 @@ bool convert_with_components(const char* model_path,
         model_loader.convert_tensors_name();
     }
 
-    return export_loaded_model(model_loader, output_path, output_type, tensor_type_rules);
+    return export_loaded_model(model_loader, output_path, output_type, tensor_type_rules, n_threads);
 }
 
 bool convert(const char* input_path,
@@ -179,5 +401,6 @@ bool convert(const char* input_path,
                                    output_path,
                                    output_type,
                                    tensor_type_rules,
-                                   convert_name);
+                                   convert_name,
+                                   0);
 }

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__

src/model_loader.cpp

@@ -943,7 +943,8 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
 
 bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                                bool enable_mmap,
-                               const std::set<std::string>* target_tensor_names) {
+                               const std::set<std::string>* target_tensor_names,
+                               bool log_progress) {
     process_model_files(enable_mmap, false);
 
     std::atomic<int64_t> read_time_ms(0);
@@ -1212,7 +1213,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
             }
             size_t curr_num       = total_tensors_processed + current_idx;
             float elapsed_seconds = (ggml_time_ms() - t_start) / 1000.0f;
-            if (total_tensors_to_process > 0) {
+            if (log_progress && total_tensors_to_process > 0) {
                 pretty_bytes_progress(static_cast<int>(curr_num),
                                       static_cast<int>(total_tensors_to_process),
                                       bytes_processed.load(),
@@ -1230,27 +1231,81 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
             break;
         }
         total_tensors_processed += tensors_to_process.size();
-        if (total_tensors_to_process > 0) {
+        if (log_progress && total_tensors_to_process > 0) {
             pretty_bytes_progress(static_cast<int>(total_tensors_processed),
                                   static_cast<int>(total_tensors_to_process),
                                   bytes_processed.load(),
                                   (ggml_time_ms() - t_start) / 1000.0f);
         }
-        if (total_tensors_processed < total_tensors_to_process && total_tensors_to_process > 0) {
+        if (log_progress && total_tensors_processed < total_tensors_to_process && total_tensors_to_process > 0) {
             printf("\n");
         }
     }
 
     int64_t end_time = ggml_time_ms();
-    LOG_INFO("loading tensors completed, taking %.2fs (read: %.2fs, memcpy: %.2fs, convert: %.2fs, copy_to_backend: %.2fs)",
-             (end_time - start_time) / 1000.f,
-             (read_time_ms.load() / (float)last_n_threads) / 1000.f,
-             (memcpy_time_ms.load() / (float)last_n_threads) / 1000.f,
-             (convert_time_ms.load() / (float)last_n_threads) / 1000.f,
-             (copy_to_backend_time_ms.load() / (float)last_n_threads) / 1000.f);
+    if (log_progress) {
+        LOG_INFO("loading tensors completed, taking %.2fs (read: %.2fs, memcpy: %.2fs, convert: %.2fs, copy_to_backend: %.2fs)",
+                 (end_time - start_time) / 1000.f,
+                 (read_time_ms.load() / (float)last_n_threads) / 1000.f,
+                 (memcpy_time_ms.load() / (float)last_n_threads) / 1000.f,
+                 (convert_time_ms.load() / (float)last_n_threads) / 1000.f,
+                 (copy_to_backend_time_ms.load() / (float)last_n_threads) / 1000.f);
+    }
     return success;
 }
 
+bool ModelLoader::load_tensor(const TensorStorage& tensor_storage, ggml_tensor* dst_tensor) {
+    if (dst_tensor == nullptr || dst_tensor->data == nullptr) {
+        LOG_ERROR("load tensor failed: null destination for '%s'", tensor_storage.name.c_str());
+        return false;
+    }
+
+    bool loaded = false;
+    std::set<std::string> target_tensor_names{tensor_storage.name};
+    auto on_new_tensor_cb = [&](const TensorStorage& current_tensor_storage, ggml_tensor** out_tensor) -> bool {
+        *out_tensor = nullptr;
+        if (current_tensor_storage.name != tensor_storage.name) {
+            return true;
+        }
+
+        if (current_tensor_storage.file_index != tensor_storage.file_index ||
+            current_tensor_storage.offset != tensor_storage.offset ||
+            current_tensor_storage.index_in_zip != tensor_storage.index_in_zip) {
+            LOG_ERROR("load tensor failed: storage mismatch for '%s'", tensor_storage.name.c_str());
+            return false;
+        }
+
+        if (current_tensor_storage.n_dims != tensor_storage.n_dims ||
+            current_tensor_storage.nelements() != tensor_storage.nelements()) {
+            LOG_ERROR("load tensor failed: metadata changed for '%s'", tensor_storage.name.c_str());
+            return false;
+        }
+
+        for (int i = 0; i < current_tensor_storage.n_dims; i++) {
+            if (current_tensor_storage.ne[i] != dst_tensor->ne[i]) {
+                LOG_ERROR("load tensor failed: shape mismatch for '%s'", tensor_storage.name.c_str());
+                return false;
+            }
+        }
+
+        *out_tensor = dst_tensor;
+        loaded      = true;
+        return true;
+    };
+
+    if (!load_tensors(on_new_tensor_cb, false, &target_tensor_names, false)) {
+        LOG_ERROR("load tensor failed: '%s'", tensor_storage.name.c_str());
+        return false;
+    }
+
+    if (!loaded) {
+        LOG_ERROR("load tensor failed: tensor '%s' not found", tensor_storage.name.c_str());
+        return false;
+    }
+
+    return true;
+}
+
 bool ModelLoader::load_float_tensor(const std::string& name,
                                     std::vector<float>& data,
                                     int n_threads,

src/model_loader.h

@@ -70,7 +70,8 @@ public:
                                               bool writable                        = true);
     bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                       bool use_mmap                                    = false,
-                      const std::set<std::string>* target_tensor_names = nullptr);
+                      const std::set<std::string>* target_tensor_names = nullptr,
+                      bool log_progress                                = true);
     bool load_tensors(std::map<std::string, ggml_tensor*>& tensors,
                       std::set<std::string> ignore_tensors = {},
                       bool use_mmap                        = false);
@@ -78,6 +79,7 @@ public:
                            std::vector<float>& data,
                            int n_threads = 0,
                            bool use_mmap = false);
+    bool load_tensor(const TensorStorage& tensor_storage, ggml_tensor* dst_tensor);
 
     std::vector<std::string> get_tensor_names() const {
         std::vector<std::string> names;