159 lines
8.8 KiB
C++
159 lines
8.8 KiB
C++
#include "gguf.hpp"
|
|
#include "skeleton.hpp"
|
|
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <fstream>
|
|
#include <limits>
|
|
#include <string>
|
|
|
|
namespace kimodo::detail {
|
|
namespace {
|
|
constexpr std::uint32_t gguf_magic = 0x46554747; // "GGUF", little endian
|
|
constexpr std::uint32_t max_metadata = 100000;
|
|
constexpr std::uint64_t max_string_bytes = 16 * 1024 * 1024;
|
|
|
|
template <class T> bool read(std::istream &in, T &value) {
|
|
return static_cast<bool>(in.read(reinterpret_cast<char *>(&value), sizeof(value)));
|
|
}
|
|
bool read_string(std::istream &in, std::string &out) {
|
|
std::uint64_t size = 0;
|
|
if (!read(in, size) || size > max_string_bytes) return false;
|
|
out.resize(static_cast<size_t>(size));
|
|
return size == 0 || static_cast<bool>(in.read(out.data(), static_cast<std::streamsize>(size)));
|
|
}
|
|
bool skip(std::istream &in, std::uint32_t type) {
|
|
// GGUF metadata scalar types. Arrays are rejected below: Kimodo metadata
|
|
// never needs them in this small, hostile-input loader.
|
|
static constexpr std::array<unsigned, 12> widths{1, 1, 2, 2, 4, 4, 4, 1, 1, 1, 8, 8};
|
|
if (type >= widths.size()) return false;
|
|
in.seekg(widths[type], std::ios::cur);
|
|
return static_cast<bool>(in);
|
|
}
|
|
}
|
|
|
|
std::expected<gguf_file, std::string> read_gguf_header(std::string_view path) {
|
|
std::ifstream in(std::string(path), std::ios::binary);
|
|
if (!in) return std::unexpected("cannot open GGUF file");
|
|
std::uint32_t magic = 0, version = 0;
|
|
std::uint64_t tensor_count = 0, metadata_count = 0;
|
|
if (!read(in, magic) || !read(in, version) || !read(in, tensor_count) || !read(in, metadata_count))
|
|
return std::unexpected("truncated GGUF header");
|
|
if (magic != gguf_magic) return std::unexpected("not a GGUF file");
|
|
if (version < 2 || version > 3) return std::unexpected("unsupported GGUF version");
|
|
if (metadata_count > max_metadata || tensor_count > 10000000)
|
|
return std::unexpected("GGUF count exceeds safety limit");
|
|
gguf_file result;
|
|
for (std::uint64_t i = 0; i < metadata_count; ++i) {
|
|
std::string key;
|
|
std::uint32_t type = 0;
|
|
if (!read_string(in, key) || !read(in, type)) return std::unexpected("truncated GGUF metadata");
|
|
if (key.empty() || key.size() > 4096) return std::unexpected("invalid GGUF metadata key");
|
|
if (type == 8) {
|
|
std::string value;
|
|
if (!read_string(in, value)) return std::unexpected("invalid GGUF string metadata");
|
|
result.strings.emplace(std::move(key), std::move(value));
|
|
} else if (type == 10 || type == 11) {
|
|
std::uint64_t value = 0;
|
|
if (!read(in, value)) return std::unexpected("truncated GGUF integer metadata");
|
|
result.uints.emplace(std::move(key), value);
|
|
} else if (type == 9) {
|
|
return std::unexpected("GGUF array metadata is not accepted by the minimal loader");
|
|
} else if (!skip(in, type)) {
|
|
return std::unexpected("unsupported or truncated GGUF metadata");
|
|
}
|
|
}
|
|
// Tensor directory is also untrusted. Read every descriptor even though
|
|
// tensor mapping/graph construction is deferred, preventing a truncated
|
|
// or malformed file from being accepted merely because its metadata is
|
|
// well formed.
|
|
std::uint64_t max_tensor_end = 0;
|
|
for (std::uint64_t i = 0; i < tensor_count; ++i) {
|
|
std::string name;
|
|
std::uint32_t dimensions = 0, type = 0;
|
|
if (!read_string(in, name) || !read(in, dimensions) || name.empty() || dimensions == 0 || dimensions > 4)
|
|
return std::unexpected("invalid GGUF tensor descriptor");
|
|
if (!result.tensor_names.emplace(name).second) return std::unexpected("duplicate GGUF tensor name");
|
|
std::uint64_t elements = 1;
|
|
for (std::uint32_t dim = 0; dim < dimensions; ++dim) {
|
|
std::uint64_t extent = 0;
|
|
if (!read(in, extent) || extent == 0 || extent > 100000000 || elements > std::numeric_limits<std::uint64_t>::max() / extent)
|
|
return std::unexpected("invalid GGUF tensor dimension");
|
|
elements *= extent;
|
|
}
|
|
std::uint64_t offset = 0;
|
|
if (!read(in, type) || !read(in, offset) || type != 0 || offset % 32 != 0 || elements > std::numeric_limits<std::uint64_t>::max() / 4)
|
|
return std::unexpected("unsupported GGUF tensor type or offset");
|
|
if (offset > std::numeric_limits<std::uint64_t>::max() - elements * 4)
|
|
return std::unexpected("GGUF tensor offset overflows");
|
|
max_tensor_end = std::max(max_tensor_end, offset + elements * 4);
|
|
}
|
|
const auto directory_end = static_cast<std::uint64_t>(in.tellg());
|
|
if (directory_end == std::numeric_limits<std::uint64_t>::max()) return std::unexpected("invalid GGUF tensor directory");
|
|
const auto data_start = (directory_end + 31U) & ~std::uint64_t{31U};
|
|
in.seekg(0, std::ios::end);
|
|
const auto file_end = static_cast<std::uint64_t>(in.tellg());
|
|
if (file_end < data_start || max_tensor_end > file_end - data_start)
|
|
return std::unexpected("GGUF tensor data is truncated");
|
|
result.tensor_count = tensor_count;
|
|
return result;
|
|
}
|
|
|
|
std::expected<void, std::string> validate_motion_gguf(const gguf_file &file) {
|
|
const auto architecture = file.strings.find("general.architecture");
|
|
if (architecture == file.strings.end() || architecture->second != "kimodo-motion")
|
|
return std::unexpected("GGUF is not a Kimodo motion model");
|
|
const auto format = file.uints.find("kimodo.format_version");
|
|
if (format == file.uints.end() || format->second != 1)
|
|
return std::unexpected("unsupported Kimodo motion GGUF format");
|
|
const auto skeleton = file.strings.find("kimodo.skeleton");
|
|
if (skeleton == file.strings.end() || !find_skeleton(skeleton->second))
|
|
return std::unexpected("motion GGUF has an unsupported skeleton");
|
|
const auto &spec = *find_skeleton(skeleton->second);
|
|
const auto motion_dim = file.uints.find("kimodo.motion_dim");
|
|
const auto body_dim = file.uints.find("kimodo.body_dim");
|
|
if (motion_dim == file.uints.end() || motion_dim->second != spec.motion_dim() ||
|
|
body_dim == file.uints.end() || body_dim->second != spec.body_dim())
|
|
return std::unexpected("motion GGUF dimensions do not match its skeleton");
|
|
const auto width = file.uints.find("kimodo.text_embedding_width");
|
|
if (width == file.uints.end() || width->second != 4096)
|
|
return std::unexpected("motion GGUF has incompatible text embedding width");
|
|
if (file.tensor_count != 414)
|
|
return std::unexpected("motion GGUF has an unexpected tensor count");
|
|
for (const char *stage : {"root_model", "body_model"}) {
|
|
for (const char *name : {"embed_text.weight", "embed_text.bias", "input_linear.weight", "input_linear.bias",
|
|
"output_linear.weight", "output_linear.bias", "linear_first_heading_angle.weight",
|
|
"linear_first_heading_angle.bias", "embed_timestep.time_embed.0.weight",
|
|
"embed_timestep.time_embed.0.bias", "embed_timestep.time_embed.2.weight",
|
|
"embed_timestep.time_embed.2.bias"}) {
|
|
if (!file.tensor_names.contains(std::string(stage) + "." + name))
|
|
return std::unexpected("motion GGUF is missing a transformer projection tensor");
|
|
}
|
|
for (unsigned layer = 0; layer < 16; ++layer) {
|
|
const std::string prefix = std::string(stage) + ".seqTransEncoder.layers." + std::to_string(layer) + ".";
|
|
for (const char *name : {"linear1.weight", "linear1.bias", "linear2.weight", "linear2.bias", "norm1.weight",
|
|
"norm1.bias", "norm2.weight", "norm2.bias", "self_attn.in_proj_weight",
|
|
"self_attn.in_proj_bias", "self_attn.out_proj.weight", "self_attn.out_proj.bias"}) {
|
|
if (!file.tensor_names.contains(prefix + name)) return std::unexpected("motion GGUF is missing a transformer layer tensor");
|
|
}
|
|
}
|
|
}
|
|
for (const char *group : {"global_root", "local_root", "body"})
|
|
for (const char *stat : {"mean", "std"})
|
|
if (!file.tensor_names.contains(std::string("stats.") + group + "." + stat))
|
|
return std::unexpected("motion GGUF is missing normalization statistics");
|
|
return {};
|
|
}
|
|
|
|
std::expected<void, std::string> validate_text_gguf(const gguf_file &file) {
|
|
const auto architecture = file.strings.find("general.architecture");
|
|
if (architecture == file.strings.end() || architecture->second != "kimodo-llm2vec")
|
|
return std::unexpected("GGUF is not a Kimodo LLM2Vec model");
|
|
const auto width = file.uints.find("kimodo.text_embedding_width");
|
|
if (width == file.uints.end() || width->second != 4096)
|
|
return std::unexpected("text GGUF has incompatible embedding width");
|
|
return {};
|
|
}
|
|
|
|
} // namespace kimodo::detail
|