stable-diffusion.cpp/src/model_io/safetensors_io.cpp

#include "safetensors_io.h"

#include <algorithm>
#include <cstdint>
#include <exception>
#include <filesystem>
#include <fstream>
#include <ostream>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include "binary_io.h"
#include "core/util.h"
#include "json.hpp"

namespace fs = std::filesystem;

static constexpr size_t ST_HEADER_SIZE_LEN = 8;

static void set_error(std::string* error, const std::string& message) {
    if (error != nullptr) {
        *error = message;
    }
}

static std::string resolve_index_shard_path(const std::string& index_path, const std::string& shard_path) {
    fs::path shard_fs_path(shard_path);
    if (shard_fs_path.is_absolute()) {
        return shard_fs_path.lexically_normal().string();
    }
    return (fs::path(index_path).parent_path() / shard_fs_path).lexically_normal().string();
}

bool is_safetensors_file(const std::string& file_path) {
    std::ifstream file(file_path, std::ios::binary);
    if (!file.is_open()) {
        return false;
    }

    // get file size
    file.seekg(0, file.end);
    size_t file_size_ = file.tellg();
    file.seekg(0, file.beg);

    // read header size
    if (file_size_ <= ST_HEADER_SIZE_LEN) {
        return false;
    }

    uint8_t header_size_buf[ST_HEADER_SIZE_LEN];
    file.read((char*)header_size_buf, ST_HEADER_SIZE_LEN);
    if (!file) {
        return false;
    }

    size_t header_size_ = model_io::read_u64(header_size_buf);
    if (header_size_ > file_size_ - ST_HEADER_SIZE_LEN || header_size_ <= 2) {
        return false;
    }

    // read header
    std::vector<char> header_buf;
    header_buf.resize(header_size_ + 1);
    header_buf[header_size_] = '\0';
    file.read(header_buf.data(), header_size_);
    if (!file) {
        return false;
    }
    try {
        nlohmann::json header_ = nlohmann::json::parse(header_buf.data());
    } catch (const std::exception&) {
        return false;
    }
    return true;
}

static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) {
    ggml_type ttype = GGML_TYPE_COUNT;
    if (dtype == "F16") {
        ttype = GGML_TYPE_F16;
    } else if (dtype == "BF16") {
        ttype = GGML_TYPE_BF16;
    } else if (dtype == "F32") {
        ttype = GGML_TYPE_F32;
    } else if (dtype == "F64") {
        ttype = GGML_TYPE_F32;
    } else if (dtype == "F8_E4M3") {
        ttype = GGML_TYPE_F16;
    } else if (dtype == "F8_E5M2") {
        ttype = GGML_TYPE_F16;
    } else if (dtype == "I32") {
        ttype = GGML_TYPE_I32;
    } else if (dtype == "I64") {
        ttype = GGML_TYPE_I32;
    } else if (dtype == "I8") {
        ttype = GGML_TYPE_I8;
    }
    return ttype;
}

struct ComfyQuantConfig {
    std::string format;
    bool convrot   = false;
    int group_size = 0;
};

static bool read_comfy_quant_config(std::ifstream& file,
                                    const std::string& file_path,
                                    const std::string& tensor_name,
                                    size_t offset,
                                    size_t size,
                                    ComfyQuantConfig& config,
                                    std::string* error) {
    static constexpr size_t MAX_COMFY_QUANT_CONFIG_SIZE = 64 * 1024;
    if (size == 0 || size > MAX_COMFY_QUANT_CONFIG_SIZE) {
        set_error(error, "invalid ComfyUI quantization metadata tensor '" + tensor_name + "' in '" + file_path + "'");
        return false;
    }

    std::vector<char> data(size + 1, '\0');
    file.clear();
    file.seekg((std::streamoff)offset, std::ios::beg);
    file.read(data.data(), (std::streamsize)size);
    if (!file) {
        set_error(error, "read ComfyUI quantization metadata tensor failed: '" + tensor_name + "'");
        return false;
    }

    try {
        const nlohmann::json json = nlohmann::json::parse(data.data(), data.data() + size);
        config.format             = json.value("format", "");
        config.convrot            = json.value("convrot", false);
        config.group_size         = json.value("convrot_groupsize", 0);
    } catch (const std::exception&) {
        set_error(error, "parsing ComfyUI quantization metadata tensor failed: '" + tensor_name + "'");
        return false;
    }
    return true;
}

// https://huggingface.co/docs/safetensors/index
bool read_safetensors_file(const std::string& file_path,
                           std::vector<TensorStorage>& tensor_storages,
                           std::string* error,
                           std::map<std::string, std::string>* metadata) {
    std::ifstream file(file_path, std::ios::binary);
    if (!file.is_open()) {
        set_error(error, "failed to open '" + file_path + "'");
        return false;
    }

    // get file size
    file.seekg(0, file.end);
    size_t file_size_ = file.tellg();
    file.seekg(0, file.beg);

    // read header size
    if (file_size_ <= ST_HEADER_SIZE_LEN) {
        set_error(error, "invalid safetensor file '" + file_path + "'");
        return false;
    }

    uint8_t header_size_buf[ST_HEADER_SIZE_LEN];
    file.read((char*)header_size_buf, ST_HEADER_SIZE_LEN);
    if (!file) {
        set_error(error, "read safetensors header size failed: '" + file_path + "'");
        return false;
    }

    size_t header_size_ = model_io::read_u64(header_size_buf);
    if (header_size_ > file_size_ - ST_HEADER_SIZE_LEN) {
        set_error(error, "invalid safetensor file '" + file_path + "'");
        return false;
    }
    const size_t data_start = ST_HEADER_SIZE_LEN + header_size_;

    // read header
    std::vector<char> header_buf;
    header_buf.resize(header_size_ + 1);
    header_buf[header_size_] = '\0';
    file.read(header_buf.data(), header_size_);
    if (!file) {
        set_error(error, "read safetensors header failed: '" + file_path + "'");
        return false;
    }

    nlohmann::json header_;
    try {
        header_ = nlohmann::json::parse(header_buf.data());
    } catch (const std::exception&) {
        set_error(error, "parsing safetensors header failed: '" + file_path + "'");
        return false;
    }

    if (metadata != nullptr) {
        metadata->clear();
        auto metadata_item = header_.find("__metadata__");
        if (metadata_item != header_.end() && metadata_item->is_object()) {
            for (const auto& item : metadata_item->items()) {
                if (item.value().is_string()) {
                    metadata->emplace(item.key(), item.value().get<std::string>());
                }
            }
        }
    }

    std::unordered_map<std::string, ComfyQuantConfig> comfy_quant_configs;
    for (const auto& item : header_.items()) {
        const std::string& name = item.key();
        if (name == "__metadata__" || !ends_with(name, ".comfy_quant")) {
            continue;
        }

        const nlohmann::json& tensor_info = item.value();
        if (tensor_info.value("dtype", "") != "U8") {
            set_error(error, "invalid dtype for ComfyUI quantization metadata tensor '" + name + "'");
            return false;
        }
        const size_t begin = tensor_info["data_offsets"][0].get<size_t>();
        const size_t end   = tensor_info["data_offsets"][1].get<size_t>();
        if (begin > end || end > file_size_ - data_start) {
            set_error(error, "data offsets out of bounds for tensor '" + name + "'");
            return false;
        }

        ComfyQuantConfig config;
        if (!read_comfy_quant_config(file, file_path, name, data_start + begin, end - begin, config, error)) {
            return false;
        }
        const std::string module_name = name.substr(0, name.size() - std::string(".comfy_quant").size());
        comfy_quant_configs.emplace(module_name, std::move(config));
    }

    tensor_storages.clear();
    for (auto& item : header_.items()) {
        std::string name           = item.key();
        nlohmann::json tensor_info = item.value();
        // LOG_DEBUG("%s %s\n", name.c_str(), tensor_info.dump().c_str());

        if (name == "__metadata__") {
            continue;
        }

        std::string dtype    = tensor_info["dtype"];
        nlohmann::json shape = tensor_info["shape"];

        if (dtype == "U8") {
            continue;
        }

        size_t begin = tensor_info["data_offsets"][0].get<size_t>();
        size_t end   = tensor_info["data_offsets"][1].get<size_t>();
        if (begin > end || end > file_size_ - data_start) {
            set_error(error, "data offsets out of bounds for tensor '" + name + "'");
            return false;
        }

        ggml_type type = safetensors_dtype_to_ggml_type(dtype);
        if (type == GGML_TYPE_COUNT) {
            set_error(error, "unsupported dtype '" + dtype + "' (tensor '" + name + "')");
            return false;
        }

        if (shape.size() > SD_MAX_DIMS) {
            set_error(error, "invalid tensor '" + name + "'");
            return false;
        }

        int n_dims              = (int)shape.size();
        int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
        for (int i = 0; i < n_dims; i++) {
            ne[i] = shape[i].get<int64_t>();
        }

        if (n_dims == 5) {
            n_dims = 4;
            ne[0]  = ne[0] * ne[1];
            ne[1]  = ne[2];
            ne[2]  = ne[3];
            ne[3]  = ne[4];
        }

        // ggml_n_dims returns 1 for scalars
        if (n_dims == 0) {
            n_dims = 1;
        }

        TensorStorage tensor_storage(name, type, ne, n_dims, 0, data_start + begin);
        tensor_storage.reverse_ne();

        if (ends_with(name, ".weight")) {
            const std::string module_name = name.substr(0, name.size() - std::string(".weight").size());
            auto config                   = comfy_quant_configs.find(module_name);
            if (config != comfy_quant_configs.end() && config->second.format == "int8_tensorwise") {
                if (type != GGML_TYPE_I8) {
                    set_error(error, "ComfyUI int8_tensorwise weight is not I8: '" + name + "'");
                    return false;
                }
                if (config->second.convrot) {
                    int group_size_remainder = config->second.group_size;
                    while (group_size_remainder > 1 && group_size_remainder % 4 == 0) {
                        group_size_remainder /= 4;
                    }
                    if (group_size_remainder != 1 || tensor_storage.ne[0] % config->second.group_size != 0) {
                        set_error(error, "invalid ComfyUI convrot group size for tensor '" + name + "'");
                        return false;
                    }
                }
                tensor_storage.is_int8_tensorwise      = true;
                tensor_storage.int8_convrot            = config->second.convrot;
                tensor_storage.int8_convrot_group_size = config->second.group_size;
            }
        } else if (ends_with(name, ".weight_scale")) {
            const std::string module_name = name.substr(0, name.size() - std::string(".weight_scale").size());
            auto config                   = comfy_quant_configs.find(module_name);
            if (config != comfy_quant_configs.end() && config->second.format == "int8_tensorwise" &&
                tensor_storage.n_dims == 2 && tensor_storage.ne[0] == 1) {
                tensor_storage.ne[0]  = tensor_storage.ne[1];
                tensor_storage.ne[1]  = 1;
                tensor_storage.n_dims = 1;
            }
        }

        size_t tensor_data_size = end - begin;

        bool tensor_size_ok;
        if (dtype == "F8_E4M3") {
            tensor_storage.is_f8_e4m3 = true;
            // f8 -> f16
            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
        } else if (dtype == "F8_E5M2") {
            tensor_storage.is_f8_e5m2 = true;
            // f8 -> f16
            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
        } else if (dtype == "F64") {
            tensor_storage.is_f64 = true;
            // f64 -> f32
            tensor_size_ok = (tensor_storage.nbytes() * 2 == tensor_data_size);
        } else if (dtype == "I64") {
            tensor_storage.is_i64 = true;
            // i64 -> i32
            tensor_size_ok = (tensor_storage.nbytes() * 2 == tensor_data_size);
        } else {
            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size);
        }
        if (!tensor_size_ok) {
            set_error(error, "size mismatch for tensor '" + name + "' (" + dtype + ")");
            return false;
        }

        tensor_storages.push_back(tensor_storage);

        // LOG_DEBUG("%s %s", tensor_storage.to_string().c_str(), dtype.c_str());
    }

    return true;
}

bool read_safetensors_index_file(const std::string& file_path,
                                 std::vector<std::string>& shard_paths,
                                 std::string* error) {
    shard_paths.clear();

    std::ifstream file(file_path);
    if (!file.is_open()) {
        set_error(error, "failed to open '" + file_path + "'");
        return false;
    }

    nlohmann::json index;
    try {
        index = nlohmann::json::parse(file);
    } catch (const std::exception&) {
        set_error(error, "parsing safetensors index failed: '" + file_path + "'");
        return false;
    }

    if (!index.is_object() || !index.contains("weight_map") || !index["weight_map"].is_object()) {
        set_error(error, "invalid safetensors index '" + file_path + "'");
        return false;
    }

    std::unordered_set<std::string> seen_shard_paths;
    for (const auto& item : index["weight_map"].items()) {
        if (!item.value().is_string()) {
            set_error(error, "invalid shard path for tensor '" + item.key() + "'");
            return false;
        }

        std::string shard_path = resolve_index_shard_path(file_path,
                                                          item.value().get<std::string>());
        if (seen_shard_paths.insert(shard_path).second) {
            shard_paths.push_back(std::move(shard_path));
        }
    }

    if (shard_paths.empty()) {
        set_error(error, "safetensors index has no tensors: '" + file_path + "'");
        return false;
    }

    return true;
}

static bool ggml_type_to_safetensors_dtype(ggml_type type, std::string* dtype) {
    switch (type) {
        case GGML_TYPE_F16:
            *dtype = "F16";
            return true;
        case GGML_TYPE_BF16:
            *dtype = "BF16";
            return true;
        case GGML_TYPE_F32:
            *dtype = "F32";
            return true;
        case GGML_TYPE_I32:
            *dtype = "I32";
            return true;
        default:
            return false;
    }
}

bool write_safetensors_file(const std::string& file_path,
                            const std::vector<TensorWriteInfo>& tensors,
                            std::string* error) {
    nlohmann::ordered_json header = nlohmann::ordered_json::object();

    uint64_t data_offset = 0;
    for (const TensorWriteInfo& write_tensor : tensors) {
        ggml_tensor* tensor = write_tensor.tensor;
        if (tensor == nullptr) {
            set_error(error, "null tensor cannot be written to safetensors");
            return false;
        }

        const std::string name = ggml_get_name(tensor);
        std::string dtype;
        if (!ggml_type_to_safetensors_dtype(tensor->type, &dtype)) {
            set_error(error,
                      "unsupported safetensors dtype '" + std::string(ggml_type_name(tensor->type)) +
                          "' for tensor '" + name + "'");
            return false;
        }

        const uint64_t tensor_nbytes = ggml_nbytes(tensor);

        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 i = 0; i < write_tensor.n_dims; ++i) {
            shape.push_back(write_tensor.ne[write_tensor.n_dims - 1 - i]);
        }
        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 + tensor_nbytes);
        json_tensor_info["data_offsets"] = data_offsets;

        header[name] = json_tensor_info;
        data_offset += tensor_nbytes;
    }

    const std::string header_str = header.dump();

    std::ofstream file(file_path, std::ios::binary);
    if (!file.is_open()) {
        set_error(error, "failed to open '" + file_path + "' for writing");
        return false;
    }

    LOG_INFO("trying to save tensors to %s", file_path.c_str());
    model_io::write_u64(file, header_str.size());
    file.write(header_str.data(), header_str.size());
    if (!file) {
        set_error(error, "failed to write safetensors header to '" + file_path + "'");
        return false;
    }

    for (const TensorWriteInfo& write_tensor : tensors) {
        ggml_tensor* tensor        = write_tensor.tensor;
        const std::string name     = ggml_get_name(tensor);
        const size_t tensor_nbytes = ggml_nbytes(tensor);
        file.write((const char*)tensor->data, tensor_nbytes);
        if (!file) {
            set_error(error,
                      "failed to write tensor '" + name + "' to '" + file_path + "'");
            return false;
        }
    }

    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_;
}