diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 7892d5213..29cb391b6 100644 --- a/examples/cli/main.cpp +++ b/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(), diff --git a/examples/server/routes_sdapi.cpp b/examples/server/routes_sdapi.cpp index b87b57105..f5275f3c0 100644 --- a/examples/server/routes_sdapi.cpp +++ b/examples/server/routes_sdapi.cpp @@ -279,7 +279,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) { @@ -361,18 +361,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"); diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 5eacdddfa..0345608e5 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -520,7 +520,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, diff --git a/src/convert.cpp b/src/convert.cpp index 0b7fe2cfb..8e94a940c 100644 --- a/src/convert.cpp +++ b/src/convert.cpp @@ -1,14 +1,33 @@ +#include +#include +#include #include +#include +#include +#include #include #include +#include +#include #include +#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 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& 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); - - std::lock_guard 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"); - return false; - } - ggml_set_name(tensor, name.c_str()); +static bool collect_tensors_for_export(ModelLoader& model_loader, + ggml_type type, + const TensorTypeRules& tensor_type_rules, + std::vector& 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; +} + +static size_t export_tensor_nbytes(const TensorExportInfo& info) { + TensorStorage output_storage = info.storage; + output_storage.type = info.type; + return static_cast(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 tensor_write_plans_from_export_infos(const std::vector& tensors) { + std::vector plans; + plans.reserve(tensors.size()); + for (const TensorExportInfo& info : tensors) { + plans.push_back(tensor_write_plan_from_export_info(info)); + } + return plans; +} - 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); +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; + } - 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]; + // This portable fallback sets the final file size. A platform-specific + // posix_fallocate/ftruncate path can replace it later. + file.seekp(static_cast(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(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& 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; - *dst_tensor = tensor; - tensors.push_back(std::move(write_info)); + 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 workers; + workers.reserve(n_threads); + + auto reserve_memory = [&](size_t bytes) -> bool { + std::unique_lock lock(work_mutex); + memory_cv.wait(lock, [&]() { + return failed || reserved_bytes == 0 || reserved_bytes + bytes <= memory_budget; + }); + if (failed) { + return false; + } + 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 lock(work_mutex); + reserved_bytes -= std::min(reserved_bytes, bytes); + } + memory_cv.notify_all(); + }; + + auto fail = [&](const std::string& message) { + { + std::lock_guard 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 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 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(tensors_written), + static_cast(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& tensors, + StreamingModelWriter& writer, + int n_threads, + std::string* error) { + std::vector 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 tensors; - bool success = load_tensors_for_export(model_loader, ggml_ctx, type, type_rules, tensors); - ggml_backend_free(backend); - + std::vector tensors; + bool success = collect_tensors_for_export(model_loader, type, type_rules, tensors); std::string error; if (success) { + std::unique_ptr writer; if (output_is_safetensors) { - success = write_safetensors_file(output_path, tensors, &error); + writer = std::make_unique(); } else { - success = write_gguf_file(output_path, tensors, &error); + writer = std::make_unique(); } + 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); } diff --git a/src/model_io/gguf_io.cpp b/src/model_io/gguf_io.cpp index c701d01f3..cd22312d5 100644 --- a/src/model_io/gguf_io.cpp +++ b/src/model_io/gguf_io.cpp @@ -1,7 +1,10 @@ #include "gguf_io.h" +#include #include +#include #include +#include #include #include @@ -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& 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(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(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(data), static_cast(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; + } +} diff --git a/src/model_io/gguf_io.h b/src/model_io/gguf_io.h index 81c981145..3a4ae8200 100644 --- a/src/model_io/gguf_io.h +++ b/src/model_io/gguf_io.h @@ -4,8 +4,12 @@ #include #include +#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& tensor_storages, @@ -14,4 +18,28 @@ bool write_gguf_file(const std::string& file_path, const std::vector& tensors, std::string* error = nullptr); +class GGUFStreamingWriter : public StreamingModelWriter { +public: + GGUFStreamingWriter() = default; + ~GGUFStreamingWriter(); + + bool write_metadata(const std::string& file_path, + const std::vector& 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 tensors_; + std::vector tensor_offsets_; + uint64_t file_size_ = 0; + ggml_context* meta_ctx_ = nullptr; + gguf_context* gguf_ctx_ = nullptr; +}; + #endif // __SD_MODEL_IO_GGUF_IO_H__ diff --git a/src/model_io/safetensors_io.cpp b/src/model_io/safetensors_io.cpp index 39131dbd8..df72b6ee9 100644 --- a/src/model_io/safetensors_io.cpp +++ b/src/model_io/safetensors_io.cpp @@ -1,8 +1,10 @@ #include "safetensors_io.h" +#include #include #include #include +#include #include #include @@ -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& 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(ST_HEADER_SIZE_LEN); ++i) { + header_size[i] = static_cast((header_str.size() >> (8 * i)) & 0xFF); + } + file.write(reinterpret_cast(header_size), sizeof(header_size)); + file.write(header_str.data(), static_cast(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(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(data), static_cast(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_; +} diff --git a/src/model_io/safetensors_io.h b/src/model_io/safetensors_io.h index 08a1bc1f3..b4938ee18 100644 --- a/src/model_io/safetensors_io.h +++ b/src/model_io/safetensors_io.h @@ -4,6 +4,7 @@ #include #include +#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& tensors, std::string* error = nullptr); +class SafetensorsStreamingWriter : public StreamingModelWriter { +public: + SafetensorsStreamingWriter() = default; + + bool write_metadata(const std::string& file_path, + const std::vector& 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 tensors_; + std::vector tensor_offsets_; + uint64_t data_start_ = 0; + uint64_t file_size_ = 0; +}; + #endif // __SD_MODEL_IO_SAFETENSORS_IO_H__ diff --git a/src/model_io/streaming_writer.h b/src/model_io/streaming_writer.h new file mode 100644 index 000000000..06f65c59b --- /dev/null +++ b/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 +#include +#include +#include + +#include "tensor_storage.h" + +class StreamingModelWriter { +public: + virtual ~StreamingModelWriter() = default; + + virtual bool write_metadata(const std::string& file_path, + const std::vector& 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__ diff --git a/src/model_io/tensor_storage.h b/src/model_io/tensor_storage.h index c0cf079c5..307535a53 100644 --- a/src/model_io/tensor_storage.h +++ b/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 on_new_tensor_cb_t; #endif // __SD_TENSOR_STORAGE_H__ diff --git a/src/model_loader.cpp b/src/model_loader.cpp index ba8e090d4..51091d98f 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -943,7 +943,8 @@ std::vector ModelLoader::mmap_tensors(std::map* target_tensor_names) { + const std::set* target_tensor_names, + bool log_progress) { process_model_files(enable_mmap, false); std::atomic 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(curr_num), static_cast(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(total_tensors_processed), static_cast(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 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& data, int n_threads, diff --git a/src/model_loader.h b/src/model_loader.h index 529f3e890..6c19c53c5 100644 --- a/src/model_loader.h +++ b/src/model_loader.h @@ -70,7 +70,8 @@ class ModelLoader { bool writable = true); bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, bool use_mmap = false, - const std::set* target_tensor_names = nullptr); + const std::set* target_tensor_names = nullptr, + bool log_progress = true); bool load_tensors(std::map& tensors, std::set ignore_tensors = {}, bool use_mmap = false); @@ -78,6 +79,7 @@ class ModelLoader { std::vector& data, int n_threads = 0, bool use_mmap = false); + bool load_tensor(const TensorStorage& tensor_storage, ggml_tensor* dst_tensor); std::vector get_tensor_names() const { std::vector names;