fixed int8 convrot loading

This commit is contained in:
Concedo 2026-07-13 18:39:32 +08:00
parent 91b1e1c89a
commit b368c874e9
2 changed files with 83 additions and 6 deletions

View file

@ -3,6 +3,7 @@
#include <cmath>
#include <cstdint>
#include <exception>
#include <istream>
#include <memory>
#include <string>
#include <unordered_map>
@ -105,17 +106,85 @@ inline LayerMap read_quantization_metadata(const nlohmann::json& header) {
return layers;
}
inline bool should_skip_side_tensor(const std::string& name, const LayerMap& quant_layers) {
inline bool should_skip_side_tensor(const nlohmann::json& header, const std::string& name, const LayerMap& quant_layers) {
if (!ends_with(name, ".weight_scale")) {
return false;
}
std::string base_name = name.substr(0, name.size() - std::string(".weight_scale").size());
auto layer_it = quant_layers.find(base_name);
return layer_it != quant_layers.end() && layer_it->second.format == "int8_tensorwise";
if (layer_it != quant_layers.end() && layer_it->second.format == "int8_tensorwise") {
return true;
}
std::string comfy_quant_name = base_name + ".comfy_quant";
return header.contains(comfy_quant_name) && header[comfy_quant_name].is_object();
}
inline void read_layer_fields(const nlohmann::json& layer_json, LayerInfo& layer) {
if (layer_json.contains("format") && layer_json["format"].is_string()) {
layer.format = layer_json["format"].get<std::string>();
}
if (layer_json.contains("convrot") && layer_json["convrot"].is_boolean()) {
layer.convrot = layer_json["convrot"].get<bool>();
}
if (layer_json.contains("convrot_groupsize") && layer_json["convrot_groupsize"].is_number_integer()) {
layer.convrot_groupsize = layer_json["convrot_groupsize"].get<int>();
} else if (layer_json.contains("convrot_group_size") && layer_json["convrot_group_size"].is_number_integer()) {
layer.convrot_groupsize = layer_json["convrot_group_size"].get<int>();
}
}
inline bool read_comfy_quant_layer(const nlohmann::json& header,
std::istream& file,
const std::string& base_name,
size_t data_start,
size_t file_size,
LayerInfo& layer,
std::string* error) {
std::string comfy_quant_name = base_name + ".comfy_quant";
if (!header.contains(comfy_quant_name) || !header[comfy_quant_name].is_object()) {
set_error(error, "unsupported dtype 'I8' without int8_tensorwise metadata (tensor '" + base_name + ".weight')");
return false;
}
const nlohmann::json& comfy_info = header[comfy_quant_name];
if (!comfy_info.contains("dtype") || comfy_info["dtype"].get<std::string>() != "U8") {
set_error(error, "unsupported comfy_quant dtype (tensor '" + comfy_quant_name + "')");
return false;
}
size_t begin = comfy_info["data_offsets"][0].get<size_t>();
size_t end = comfy_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 '" + comfy_quant_name + "'");
return false;
}
size_t payload_size = end - begin;
std::vector<char> payload(payload_size + 1, '\0');
std::streampos previous_pos = file.tellg();
file.seekg(data_start + begin);
file.read(payload.data(), payload_size);
if (previous_pos != std::streampos(-1)) {
file.seekg(previous_pos);
}
if (!file) {
set_error(error, "read comfy_quant metadata failed: '" + comfy_quant_name + "'");
return false;
}
try {
nlohmann::json layer_json = nlohmann::json::parse(payload.data());
read_layer_fields(layer_json, layer);
} catch (const std::exception&) {
set_error(error, "parsing comfy_quant metadata failed: '" + comfy_quant_name + "'");
return false;
}
return true;
}
inline bool fill_i8_tensorwise_storage(const nlohmann::json& header,
std::istream& file,
const LayerMap& quant_layers,
const std::string& name,
size_t data_start,
@ -130,7 +199,14 @@ inline bool fill_i8_tensorwise_storage(const nlohmann::json& header,
std::string base_name = name.substr(0, name.size() - std::string(".weight").size());
auto layer_it = quant_layers.find(base_name);
if (layer_it == quant_layers.end() || layer_it->second.format != "int8_tensorwise") {
LayerInfo layer;
if (layer_it != quant_layers.end()) {
layer = layer_it->second;
} else if (!read_comfy_quant_layer(header, file, base_name, data_start, file_size, layer, error)) {
return false;
}
if (layer.format != "int8_tensorwise") {
set_error(error, "unsupported dtype 'I8' without int8_tensorwise metadata (tensor '" + name + "')");
return false;
}
@ -183,8 +259,8 @@ inline bool fill_i8_tensorwise_storage(const nlohmann::json& header,
}
ext->is_i8_tensorwise = true;
ext->convrot = layer_it->second.convrot;
ext->convrot_groupsize = layer_it->second.convrot_groupsize;
ext->convrot = layer.convrot;
ext->convrot_groupsize = layer.convrot_groupsize;
ext->scale_offset = data_start + scale_begin;
tensor_storage.kcpp_ext = ext;
if (tensor_storage.nbytes_to_read() != (int64_t)tensor_data_size) {

View file

@ -152,7 +152,7 @@ bool read_safetensors_file(const std::string& file_path,
continue;
}
if (kcpp_safetensors_quant::should_skip_side_tensor(name, quant_layers)) {
if (kcpp_safetensors_quant::should_skip_side_tensor(header_, name, quant_layers)) {
continue;
}
@ -208,6 +208,7 @@ bool read_safetensors_file(const std::string& file_path,
bool tensor_size_ok;
if (dtype == "I8") {
if (!kcpp_safetensors_quant::fill_i8_tensorwise_storage(header_,
file,
quant_layers,
name,
data_start,