diff --git a/ggml b/ggml index f31dab0fc..0de1c777c 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit f31dab0fc95e8b2101f44cd34f8f53f6cc2cfb6c +Subproject commit 0de1c777cb20b167afd80677e4213f96f500f8d4 diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index d0683166d..e5a5a9a9b 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -34,6 +34,7 @@ namespace MiniMaxH3 { int64_t time_embed_dim = 2688; int64_t rope_inv_freq_len = 16; int64_t adaln_curve_grid = 0; + bool is_comfyui_layout = false; int patch_t = 1; int patch_h = 2; int patch_w = 2; @@ -75,6 +76,10 @@ namespace MiniMaxH3 { } if (const auto* weight = find("audio_patch_proj.weight")) { config.audio_latent_channels = weight->ne[0]; + // ComfyUI GGUF exports can store video_patch_proj transposed, + // so derive the transformer width from the unambiguous audio + // projection when it is wider than the video-derived value. + config.hidden_size = std::max(config.hidden_size, weight->ne[1]); } config.num_layers = count_blocks(tensors, prefix + ".blocks."); config.token_refiner_num_layers = count_blocks(tensors, prefix + ".token_refiner.blocks."); @@ -105,15 +110,19 @@ namespace MiniMaxH3 { if (const auto* inv_freq = find("rope.inv_freq")) { config.rope_inv_freq_len = inv_freq->ne[0]; } + const auto* adaln = find("blocks.0.adaln_proj.linear.weight"); + config.is_comfyui_layout = adaln != nullptr && adaln->has_comfy_original_shape; LOG_DEBUG("minimax_h3: layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64 - ", head_dim=%" PRId64 ", ffn=%" PRId64 ", adaln_curve=%" PRId64, + ", head_dim=%" PRId64 ", ffn=%" PRId64 ", adaln_curve=%" PRId64 + ", comfyui_layout=%d", config.num_layers, config.hidden_size, config.num_attention_heads, config.attention_head_dim, config.ffn_hidden_size, - config.adaln_curve_grid); + config.adaln_curve_grid, + config.is_comfyui_layout); return config; } }; @@ -286,7 +295,8 @@ namespace MiniMaxH3 { int expand, int modalities, bool apply_silu, - bool force_f32) + bool force_f32, + bool force_prec_f32) : hidden_size(hidden_size), expand(expand), modalities(modalities), @@ -294,7 +304,8 @@ namespace MiniMaxH3 { blocks["linear"] = std::make_shared(time_dim, hidden_size * expand * modalities, true, - force_f32); + force_f32, + force_prec_f32); } ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t_emb) { @@ -403,7 +414,10 @@ namespace MiniMaxH3 { 6, 3, !config.uses_adaln_curves(), - config.uses_adaln_curves()); + config.uses_adaln_curves() && + !config.is_comfyui_layout, + config.uses_adaln_curves() && + config.is_comfyui_layout); } ggml_tensor* forward(GGMLRunnerContext* ctx, @@ -465,7 +479,10 @@ namespace MiniMaxH3 { 2, 1, !config.uses_adaln_curves(), - config.uses_adaln_curves()); + config.uses_adaln_curves() && + !config.is_comfyui_layout, + config.uses_adaln_curves() && + config.is_comfyui_layout); blocks["video_out"] = std::make_shared(config.hidden_size, video_dim, true, true); blocks["audio_out"] = std::make_shared(config.hidden_size, config.audio_latent_channels, true, true); } diff --git a/src/model_io/gguf_io.cpp b/src/model_io/gguf_io.cpp index cd22312d5..0c1a708a3 100644 --- a/src/model_io/gguf_io.cpp +++ b/src/model_io/gguf_io.cpp @@ -1,9 +1,11 @@ #include "gguf_io.h" #include +#include #include #include #include +#include #include #include #include @@ -18,6 +20,205 @@ static void set_error(std::string* error, const std::string& message) { } } +enum class ComfyShapeResult { + NOT_FOUND, + RESTORED, + INVALID, +}; + +static bool checked_element_count(const int64_t* ne, + size_t n_dims, + uint64_t* element_count) { + uint64_t count = 1; + for (size_t i = 0; i < n_dims; ++i) { + if (ne[i] <= 0 || static_cast(ne[i]) > std::numeric_limits::max() / count) { + return false; + } + count *= static_cast(ne[i]); + } + *element_count = count; + return true; +} + +static ComfyShapeResult apply_comfy_shape_values(const int64_t* shape, + size_t shape_size, + uint64_t physical_element_count, + int64_t* ne, + int* n_dims) { + if (shape_size == 0 || shape_size > GGML_MAX_DIMS + 1 || + !std::all_of(shape, shape + shape_size, [](int64_t dim) { return dim > 0; })) { + return ComfyShapeResult::INVALID; + } + + uint64_t logical_element_count = 1; + for (size_t i = 0; i < shape_size; ++i) { + const uint64_t dim = static_cast(shape[i]); + if (dim > std::numeric_limits::max() / logical_element_count) { + return ComfyShapeResult::INVALID; + } + logical_element_count *= dim; + } + if (logical_element_count != physical_element_count) { + return ComfyShapeResult::INVALID; + } + + const size_t collapsed_dims = shape_size > GGML_MAX_DIMS + ? shape_size - GGML_MAX_DIMS + 1 + : 0; + uint64_t collapsed_dimension = 1; + for (size_t i = 0; i < collapsed_dims; ++i) { + const uint64_t dim = static_cast(shape[i]); + if (dim > std::numeric_limits::max() / collapsed_dimension) { + return ComfyShapeResult::INVALID; + } + collapsed_dimension *= dim; + } + if (collapsed_dimension > static_cast(std::numeric_limits::max())) { + return ComfyShapeResult::INVALID; + } + std::fill(ne, ne + GGML_MAX_DIMS, 1); + std::reverse_copy(shape + collapsed_dims, shape + shape_size, ne); + if (collapsed_dims > 0) { + ne[GGML_MAX_DIMS - 1] = static_cast(collapsed_dimension); + *n_dims = GGML_MAX_DIMS; + } else { + *n_dims = static_cast(shape_size); + } + return ComfyShapeResult::RESTORED; +} + +// ComfyUI-GGUF may reshape a quantized tensor to satisfy the quantizer's +// block-size requirement. The original PyTorch dimensions are preserved in +// comfy.gguf.orig_shape.; TensorStorage uses GGML's reversed +// dimension order, so restore that logical shape before model detection and +// loading while retaining the physical byte layout in the file. +static ComfyShapeResult apply_comfy_original_shape(const gguf_context* ctx, + const std::string& tensor_name, + uint64_t physical_element_count, + int64_t* ne, + int* n_dims) { + const std::string key = "comfy.gguf.orig_shape." + tensor_name; + const int64_t key_id = gguf_find_key(ctx, key.c_str()); + if (key_id < 0) { + return ComfyShapeResult::NOT_FOUND; + } + if (gguf_get_kv_type(ctx, key_id) != GGUF_TYPE_ARRAY || + gguf_get_arr_type(ctx, key_id) != GGUF_TYPE_INT32) { + return ComfyShapeResult::INVALID; + } + + const size_t shape_size = gguf_get_arr_n(ctx, key_id); + if (shape_size == 0 || shape_size > GGML_MAX_DIMS + 1) { + return ComfyShapeResult::INVALID; + } + + const int32_t* shape = static_cast(gguf_get_arr_data(ctx, key_id)); + if (shape == nullptr) { + return ComfyShapeResult::INVALID; + } + std::vector shape_values(shape, shape + shape_size); + return apply_comfy_shape_values(shape_values.data(), + shape_values.size(), + physical_element_count, + ne, + n_dims); +} + +static ggml_type comfy_expected_type(ggml_type source_type, int64_t logical_row_width) { + if (logical_row_width <= 0 || + logical_row_width % ggml_blck_size(GGML_TYPE_Q4_0) != 0 || + logical_row_width % ggml_blck_size(source_type) == 0) { + return GGML_TYPE_COUNT; + } + + switch (source_type) { + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_TQ1_0: + case GGML_TYPE_TQ2_0: + return GGML_TYPE_Q4_0; + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_Q8_0: + return GGML_TYPE_Q8_0; + default: + return GGML_TYPE_COUNT; + } +} + +static void apply_comfy_expected_type(TensorStorage& tensor_storage, + bool restored_comfy_shape, + size_t* remapped_tensors) { + if (!restored_comfy_shape) { + return; + } + tensor_storage.has_comfy_original_shape = true; + const ggml_type expected_type = + comfy_expected_type(tensor_storage.type, tensor_storage.ne[0]); + if (expected_type != GGML_TYPE_COUNT) { + tensor_storage.expected_type = expected_type; + ++*remapped_tensors; + } else if (ggml_is_quantized(tensor_storage.type) && + tensor_storage.ne[0] % ggml_blck_size(tensor_storage.type) != 0) { + LOG_WARN("ComfyUI GGUF: no quantized remap for tensor %s (source type %s, logical row %" PRId64 "); falling back to F32", + tensor_storage.name.c_str(), + ggml_type_name(tensor_storage.type), + tensor_storage.ne[0]); + } +} + +static bool checked_physical_element_count(const int64_t* ne, + size_t n_dims, + uint64_t* element_count, + std::string* error, + const std::string& tensor_name) { + if (checked_element_count(ne, n_dims, element_count)) { + return true; + } + set_error(error, "invalid or overflowing physical shape for tensor '" + tensor_name + "'"); + return false; +} + +static ComfyShapeResult apply_fallback_comfy_shape(const GGUFReader& reader, + const GGUFTensorInfo& tensor_info, + uint64_t physical_element_count, + int64_t* logical_ne, + int* logical_n_dims, + std::string* error) { + const auto* shape = reader.comfy_original_shape(tensor_info.name); + if (shape == nullptr) { + return ComfyShapeResult::NOT_FOUND; + } + const ComfyShapeResult result = apply_comfy_shape_values(shape->data(), + shape->size(), + physical_element_count, + logical_ne, + logical_n_dims); + if (result != ComfyShapeResult::RESTORED) { + if (result == ComfyShapeResult::INVALID) { + set_error(error, "invalid comfy.gguf.orig_shape for tensor '" + tensor_info.name + "'"); + } + } + return result; +} + +static void initialize_logical_shape(const std::vector& physical_shape, + int64_t* logical_ne, + int* logical_n_dims) { + std::fill(logical_ne, logical_ne + GGML_MAX_DIMS, 1); + std::copy(physical_shape.begin(), physical_shape.end(), logical_ne); + *logical_n_dims = static_cast(physical_shape.size()); +} + bool is_gguf_file(const std::string& file_path) { std::ifstream file(file_path, std::ios::binary); if (!file.is_open()) { @@ -55,31 +256,103 @@ bool read_gguf_file(const std::string& file_path, return false; } - size_t data_offset = gguf_reader.data_offset(); + size_t data_offset = gguf_reader.data_offset(); + size_t remapped_tensors = 0; for (const auto& gguf_tensor_info : gguf_reader.tensors()) { + uint64_t physical_element_count = 0; + if (!checked_physical_element_count(gguf_tensor_info.shape.data(), + gguf_tensor_info.shape.size(), + &physical_element_count, + error, + gguf_tensor_info.name)) { + return false; + } + + int64_t logical_ne[GGML_MAX_DIMS]; + int logical_n_dims = 0; + initialize_logical_shape(gguf_tensor_info.shape, logical_ne, &logical_n_dims); + const ComfyShapeResult shape_result = + apply_fallback_comfy_shape(gguf_reader, + gguf_tensor_info, + physical_element_count, + logical_ne, + &logical_n_dims, + error); + if (shape_result == ComfyShapeResult::INVALID) { + return false; + } + TensorStorage tensor_storage( gguf_tensor_info.name, gguf_tensor_info.type, - gguf_tensor_info.shape.data(), - static_cast(gguf_tensor_info.shape.size()), + logical_ne, + logical_n_dims, 0, data_offset + gguf_tensor_info.offset); + apply_comfy_expected_type(tensor_storage, + shape_result == ComfyShapeResult::RESTORED, + &remapped_tensors); tensor_storages.push_back(tensor_storage); } + if (remapped_tensors > 0) { + LOG_INFO("ComfyUI GGUF remapped %zu quantized tensors in fallback reader", remapped_tensors); + } return true; } int n_tensors = static_cast(gguf_get_n_tensors(ctx_gguf_)); - size_t data_offset = gguf_get_data_offset(ctx_gguf_); + size_t data_offset = gguf_get_data_offset(ctx_gguf_); + size_t remapped_tensors = 0; for (int i = 0; i < n_tensors; i++) { std::string name = gguf_get_tensor_name(ctx_gguf_, i); ggml_tensor* dummy = ggml_get_tensor(ctx_meta_, name.c_str()); size_t offset = data_offset + gguf_get_tensor_offset(ctx_gguf_, i); - TensorStorage tensor_storage(name, dummy->type, dummy->ne, ggml_n_dims(dummy), 0, offset); + uint64_t physical_element_count = 0; + if (!checked_physical_element_count(dummy->ne, + GGML_MAX_DIMS, + &physical_element_count, + error, + name)) { + gguf_free(ctx_gguf_); + ggml_free(ctx_meta_); + return false; + } + + int64_t logical_ne[GGML_MAX_DIMS]; + std::copy(dummy->ne, dummy->ne + GGML_MAX_DIMS, logical_ne); + int logical_n_dims = ggml_n_dims(dummy); + const ComfyShapeResult shape_result = + apply_comfy_original_shape(ctx_gguf_, + name, + physical_element_count, + logical_ne, + &logical_n_dims); + if (shape_result == ComfyShapeResult::INVALID) { + gguf_free(ctx_gguf_); + ggml_free(ctx_meta_); + set_error(error, "invalid comfy.gguf.orig_shape for tensor '" + name + "'"); + return false; + } + + TensorStorage tensor_storage(name, + dummy->type, + logical_ne, + logical_n_dims, + 0, + offset); + + // ComfyUI packs quantized tensors using a physical row width that is + // compatible with the source quant block, then records the original + // logical shape in metadata. If the logical matrix row is not valid + // for that source type, transcode it once to Q4_0 (32-value blocks) + // instead of expanding the entire matrix to F32 at runtime. + apply_comfy_expected_type(tensor_storage, + shape_result == ComfyShapeResult::RESTORED, + &remapped_tensors); if (ggml_nbytes(dummy) != tensor_storage.nbytes()) { gguf_free(ctx_gguf_); @@ -91,6 +364,10 @@ bool read_gguf_file(const std::string& file_path, tensor_storages.push_back(tensor_storage); } + if (remapped_tensors > 0) { + LOG_INFO("ComfyUI GGUF remapped %zu quantized tensors", remapped_tensors); + } + gguf_free(ctx_gguf_); ggml_free(ctx_meta_); diff --git a/src/model_io/gguf_reader_ext.h b/src/model_io/gguf_reader_ext.h index 7da20d0a5..6d0bdc252 100644 --- a/src/model_io/gguf_reader_ext.h +++ b/src/model_io/gguf_reader_ext.h @@ -1,8 +1,11 @@ #ifndef __SD_MODEL_IO_GGUF_READER_EXT_H__ #define __SD_MODEL_IO_GGUF_READER_EXT_H__ +#include #include #include +#include +#include #include #include @@ -35,6 +38,7 @@ enum class GGUFMetadataType : uint32_t { class GGUFReader { private: std::vector tensors_; + std::map> comfy_original_shapes_; size_t data_offset_; size_t alignment_ = 32; // default alignment is 32 @@ -54,6 +58,55 @@ class GGUFReader { return fin.good(); } + bool skip_value(std::ifstream& fin, uint32_t type) { + switch (static_cast(type)) { + case GGUFMetadataType::UINT8: + case GGUFMetadataType::INT8: + case GGUFMetadataType::BOOL: + return safe_seek(fin, 1, std::ios::cur); + + case GGUFMetadataType::UINT16: + case GGUFMetadataType::INT16: + return safe_seek(fin, 2, std::ios::cur); + + case GGUFMetadataType::UINT32: + case GGUFMetadataType::INT32: + case GGUFMetadataType::FLOAT32: + return safe_seek(fin, 4, std::ios::cur); + + case GGUFMetadataType::UINT64: + case GGUFMetadataType::INT64: + case GGUFMetadataType::FLOAT64: + return safe_seek(fin, 8, std::ios::cur); + + case GGUFMetadataType::STRING: { + uint64_t len = 0; + if (!safe_read(fin, len) || + len > static_cast(std::numeric_limits::max())) { + return false; + } + return safe_seek(fin, static_cast(len), std::ios::cur); + } + + case GGUFMetadataType::ARRAY: { + uint32_t elem_type = 0; + uint64_t len = 0; + if (!safe_read(fin, elem_type) || !safe_read(fin, len)) { + return false; + } + for (uint64_t i = 0; i < len; ++i) { + if (!skip_value(fin, elem_type)) { + return false; + } + } + return true; + } + + default: + return false; + } + } + bool read_metadata(std::ifstream& fin) { uint64_t key_len = 0; if (!safe_read(fin, key_len)) @@ -70,7 +123,7 @@ class GGUFReader { if (!safe_read(fin, type)) return false; - if (key == "general.alignment") { + if (key == "general.alignment" && type == static_cast(GGUFMetadataType::UINT32)) { uint32_t align_val = 0; if (!safe_read(fin, align_val)) return false; @@ -84,52 +137,43 @@ class GGUFReader { return true; } - switch (static_cast(type)) { - case GGUFMetadataType::UINT8: - case GGUFMetadataType::INT8: - case GGUFMetadataType::BOOL: - return safe_seek(fin, 1, std::ios::cur); - - case GGUFMetadataType::UINT16: - case GGUFMetadataType::INT16: - return safe_seek(fin, 2, std::ios::cur); - - case GGUFMetadataType::UINT32: - case GGUFMetadataType::INT32: - case GGUFMetadataType::FLOAT32: - return safe_seek(fin, 4, std::ios::cur); - - case GGUFMetadataType::UINT64: - case GGUFMetadataType::INT64: - case GGUFMetadataType::FLOAT64: - return safe_seek(fin, 8, std::ios::cur); - - case GGUFMetadataType::STRING: { - uint64_t len = 0; - if (!safe_read(fin, len)) - return false; - return safe_seek(fin, len, std::ios::cur); + if (type == static_cast(GGUFMetadataType::ARRAY)) { + uint32_t elem_type = 0; + uint64_t len = 0; + if (!safe_read(fin, elem_type) || !safe_read(fin, len)) { + return false; } - case GGUFMetadataType::ARRAY: { - uint32_t elem_type = 0; - uint64_t len = 0; - if (!safe_read(fin, elem_type)) - return false; - if (!safe_read(fin, len)) + constexpr const char* prefix = "comfy.gguf.orig_shape."; + if (key.compare(0, std::char_traits::length(prefix), prefix) == 0) { + if (elem_type != static_cast(GGUFMetadataType::INT32) || + len == 0 || len > GGML_MAX_DIMS + 1) { return false; + } - for (uint64_t i = 0; i < len; i++) { - if (!read_metadata(fin)) + std::vector shape; + shape.reserve(static_cast(len)); + for (uint64_t i = 0; i < len; ++i) { + int32_t dim = 0; + if (!safe_read(fin, dim)) { return false; + } + shape.push_back(dim); } + comfy_original_shapes_[key.substr(std::char_traits::length(prefix))] = + std::move(shape); return true; } - default: - LOG_ERROR("Unknown metadata type=%u", type); - return false; + for (uint64_t i = 0; i < len; ++i) { + if (!skip_value(fin, elem_type)) { + return false; + } + } + return true; } + + return skip_value(fin, type); } GGUFTensorInfo read_tensor_info(std::ifstream& fin) { @@ -155,6 +199,12 @@ class GGUFReader { if (n_dims > GGML_MAX_DIMS) { for (uint32_t i = GGML_MAX_DIMS; i < n_dims; i++) { + if (info.shape[GGML_MAX_DIMS - 1] <= 0 || + info.shape[i] <= 0 || + info.shape[GGML_MAX_DIMS - 1] > + std::numeric_limits::max() / info.shape[i]) { + throw std::runtime_error("tensor shape overflows int64"); + } info.shape[GGML_MAX_DIMS - 1] *= info.shape[i]; // stack to last dim; } info.shape.resize(GGML_MAX_DIMS); @@ -201,6 +251,7 @@ class GGUFReader { version, (unsigned long long)tensor_count, (unsigned long long)metadata_kv_count); // --- Read Metadata --- + comfy_original_shapes_.clear(); for (uint64_t i = 0; i < metadata_kv_count; i++) { if (!read_metadata(fin)) { LOG_ERROR("read meta data failed"); @@ -229,6 +280,11 @@ class GGUFReader { const std::vector& tensors() const { return tensors_; } size_t data_offset() const { return data_offset_; } + + const std::vector* comfy_original_shape(const std::string& tensor_name) const { + auto it = comfy_original_shapes_.find(tensor_name); + return it == comfy_original_shapes_.end() ? nullptr : &it->second; + } }; #endif // __SD_MODEL_IO_GGUF_READER_EXT_H__ diff --git a/src/model_io/tensor_storage.h b/src/model_io/tensor_storage.h index 307535a53..5c977f516 100644 --- a/src/model_io/tensor_storage.h +++ b/src/model_io/tensor_storage.h @@ -17,6 +17,7 @@ struct TensorStorage { std::string name; ggml_type type = GGML_TYPE_F32; ggml_type expected_type = GGML_TYPE_COUNT; + bool has_comfy_original_shape = false; bool is_f8_e4m3 = false; bool is_f8_e5m2 = false; bool is_f64 = false; diff --git a/src/model_loader.cpp b/src/model_loader.cpp index 3cbfbbd2e..a70ffefd3 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -6,7 +6,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -156,12 +158,29 @@ void convert_tensor(void* src, ggml_type src_type, void* dst, ggml_type dst_type, - int nrows, - int n_per_row, + int64_t nrows, + int64_t n_per_row, std::vector imatrix = {}) { - int n = nrows * n_per_row; + if (nrows <= 0 || n_per_row <= 0 || + nrows > std::numeric_limits::max() / n_per_row) { + throw std::runtime_error("tensor conversion dimensions overflow"); + } + const int64_t n = nrows * n_per_row; + if (ggml_is_quantized(src_type) && n % ggml_blck_size(src_type) != 0) { + throw std::runtime_error(sd_format("tensor conversion element count is not aligned for source type %s", + ggml_type_name(src_type))); + } + if (ggml_is_quantized(dst_type) && n % ggml_blck_size(dst_type) != 0) { + throw std::runtime_error(sd_format("tensor conversion element count is not aligned for destination type %s", + ggml_type_name(dst_type))); + } if (src_type == dst_type) { - size_t nbytes = n * ggml_type_size(src_type) / ggml_blck_size(src_type); + const int64_t block = ggml_blck_size(src_type); + if (n % block != 0) { + throw std::runtime_error(sd_format("tensor conversion element count is not aligned for type %s", + ggml_type_name(src_type))); + } + size_t nbytes = static_cast(n / block) * ggml_type_size(src_type); memcpy(((char*)dst), ((char*)src), nbytes); } else if (src_type == GGML_TYPE_F32) { if (dst_type == GGML_TYPE_F16) { @@ -190,16 +209,54 @@ void convert_tensor(void* src, throw std::runtime_error(sd_format("type %s unsupported for integer quantization: no dequantization available", ggml_type_name(src_type))); } - std::vector buf; - buf.resize(sizeof(float) * n); - char* src_data_f32 = buf.data(); - qtype->to_float(src, (float*)src_data_f32, n); - if (dst_type == GGML_TYPE_F16) { - ggml_fp32_to_fp16_row((float*)src_data_f32, (ggml_fp16_t*)dst, n); - } else { + // Convert in block-aligned row groups. ComfyUI may flatten a + // logical row width that is not divisible by the source quant block; + // dequantizing the full tensor would require multi-GB scratch buffers. + const int64_t src_block = ggml_blck_size(src_type); + int64_t rows_per_group = src_block / std::gcd(src_block, n_per_row); + int64_t target_rows = std::max(rows_per_group, + (1 << 20) / std::max(1, n_per_row)); + if (target_rows > std::numeric_limits::max() - rows_per_group + 1) { + throw std::runtime_error("tensor conversion row group overflows"); + } + rows_per_group = ((target_rows + rows_per_group - 1) / rows_per_group) * rows_per_group; + rows_per_group = std::min(rows_per_group, nrows); + + std::vector buf(static_cast(rows_per_group) * n_per_row); + if (dst_type != GGML_TYPE_F16) { imatrix.resize(n_per_row, 1.0f); - const float* im = imatrix.data(); - ggml_quantize_chunk(dst_type, (float*)src_data_f32, dst, 0, nrows, n_per_row, im); + } + const float* im = imatrix.empty() ? nullptr : imatrix.data(); + const size_t src_type_size = ggml_type_size(src_type); + const size_t dst_type_size = ggml_type_size(dst_type); + const int64_t dst_block = ggml_blck_size(dst_type); + + for (int64_t row = 0; row < nrows; row += rows_per_group) { + const int64_t group_rows = std::min(rows_per_group, nrows - row); + if (group_rows > std::numeric_limits::max() / n_per_row) { + throw std::runtime_error("tensor conversion group dimensions overflow"); + } + const int64_t group_elements = group_rows * n_per_row; + if (group_elements % src_block != 0) { + throw std::runtime_error(sd_format("tensor conversion source group is not aligned for type %s", + ggml_type_name(src_type))); + } + const char* src_group = static_cast(src) + + static_cast(row) * n_per_row / src_block * src_type_size; + qtype->to_float(src_group, buf.data(), group_elements); + + if (dst_type == GGML_TYPE_F16) { + auto* dst_group = static_cast(dst) + static_cast(row) * n_per_row; + ggml_fp32_to_fp16_row(buf.data(), dst_group, group_elements); + } else { + if (n_per_row % dst_block != 0) { + throw std::runtime_error(sd_format("tensor conversion destination row is not aligned for type %s", + ggml_type_name(dst_type))); + } + char* dst_group = static_cast(dst) + + static_cast(row) * n_per_row / dst_block * dst_type_size; + ggml_quantize_chunk(dst_type, buf.data(), dst_group, 0, group_rows, n_per_row, im); + } } } } @@ -1243,8 +1300,8 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, tensor_storage.type, convert_buf, dst_tensor->type, - (int)tensor_storage.nelements() / (int)tensor_storage.ne[0], - (int)tensor_storage.ne[0], + tensor_storage.nelements() / tensor_storage.ne[0], + tensor_storage.ne[0], std::move(imatrix)); } else { convert_buf = read_buf; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index fb62a7840..3673d95e6 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -36,3 +36,9 @@ target_include_directories(test-common-json PRIVATE target_link_libraries(test-common-json PRIVATE stable-diffusion zip ${CMAKE_THREAD_LIBS_INIT}) add_test(NAME test-common-json COMMAND test-common-json) + +add_executable(test-gguf-comfy-shape test-gguf-comfy-shape.cpp) +target_include_directories(test-gguf-comfy-shape PRIVATE + "${PROJECT_SOURCE_DIR}/src") +target_link_libraries(test-gguf-comfy-shape PRIVATE stable-diffusion) +add_test(NAME test-gguf-comfy-shape COMMAND test-gguf-comfy-shape) diff --git a/tests/test-gguf-comfy-shape.cpp b/tests/test-gguf-comfy-shape.cpp new file mode 100644 index 000000000..b55f45fdc --- /dev/null +++ b/tests/test-gguf-comfy-shape.cpp @@ -0,0 +1,179 @@ +#include "ggml.h" +#include "model/diffusion/minimax_h3.hpp" +#include "model_io/gguf_io.h" +#include "model_io/gguf_reader_ext.h" + +#include +#include +#include +#include +#include + +namespace { + +template +void write_value(std::ofstream& file, T value) { + file.write(reinterpret_cast(&value), sizeof(value)); +} + +void write_gguf(const std::filesystem::path& path, + const std::vector& original_shape, + ggml_type type = GGML_TYPE_F32, + const std::vector& physical_shape = {6, 2}) { + std::ofstream file(path, std::ios::binary | std::ios::trunc); + GGML_ASSERT(file.is_open()); + + file.write("GGUF", 4); + write_value(file, 3); + write_value(file, 1); + write_value(file, 1); + + const std::string metadata_key = "comfy.gguf.orig_shape.foo"; + write_value(file, metadata_key.size()); + file.write(metadata_key.data(), static_cast(metadata_key.size())); + write_value(file, 9); // ARRAY + write_value(file, 5); // INT32 + write_value(file, original_shape.size()); + for (int32_t dim : original_shape) { + write_value(file, dim); + } + + write_value(file, 3); + file.write("foo", 3); + write_value(file, physical_shape.size()); + for (uint64_t dim : physical_shape) { + write_value(file, dim); + } + write_value(file, static_cast(type)); + write_value(file, 0); + + const auto position = static_cast(file.tellp()); + const uint64_t aligned = (position + 31) & ~uint64_t(31); + for (uint64_t i = position; i < aligned; ++i) { + file.put('\0'); + } + uint64_t element_count = 1; + for (uint64_t dim : physical_shape) { + element_count *= dim; + } + const size_t data_size = + static_cast(element_count / ggml_blck_size(type)) * ggml_type_size(type); + for (size_t i = 0; i < data_size; ++i) { + file.put('\0'); + } +} + +} // namespace + +int main() { + const auto path = std::filesystem::temp_directory_path() / + "stable-diffusion-comfy-shape-test.gguf"; + + write_gguf(path, {2, 6}); + + GGUFReader fallback_reader; + GGML_ASSERT(fallback_reader.load(path.string())); + const auto* fallback_shape = fallback_reader.comfy_original_shape("foo"); + GGML_ASSERT(fallback_shape != nullptr); + GGML_ASSERT(*fallback_shape == std::vector({2, 6})); + + std::vector tensors; + std::string error; + GGML_ASSERT(read_gguf_file(path.string(), tensors, &error)); + GGML_ASSERT(tensors.size() == 1); + GGML_ASSERT(tensors[0].has_comfy_original_shape); + GGML_ASSERT(tensors[0].n_dims == 2); + GGML_ASSERT(tensors[0].ne[0] == 6); + GGML_ASSERT(tensors[0].ne[1] == 2); + + // Qwen's vision patch embedding is five-dimensional in PyTorch, so its + // leading dimensions are collapsed into GGML's four-dimensional shape. + write_gguf(path, + {1152, 3, 2, 16, 16}, + GGML_TYPE_F16, + {16, 16, 6, 1152}); + tensors.clear(); + error.clear(); + GGML_ASSERT(read_gguf_file(path.string(), tensors, &error)); + GGML_ASSERT(tensors.size() == 1); + GGML_ASSERT(tensors[0].has_comfy_original_shape); + GGML_ASSERT(tensors[0].n_dims == 4); + GGML_ASSERT(tensors[0].ne[0] == 16); + GGML_ASSERT(tensors[0].ne[1] == 16); + GGML_ASSERT(tensors[0].ne[2] == 2); + GGML_ASSERT(tensors[0].ne[3] == 3456); + + write_gguf(path, + {256, 96}, + GGML_TYPE_Q4_K, + {256, 96}); + tensors.clear(); + error.clear(); + GGML_ASSERT(read_gguf_file(path.string(), tensors, &error)); + GGML_ASSERT(tensors.size() == 1); + GGML_ASSERT(tensors[0].has_comfy_original_shape); + GGML_ASSERT(tensors[0].expected_type == GGML_TYPE_Q4_0); + GGML_ASSERT(tensors[0].ne[0] == 96); + GGML_ASSERT(tensors[0].ne[1] == 256); + + write_gguf(path, {-2, 6}); + tensors.clear(); + error.clear(); + GGML_ASSERT(!read_gguf_file(path.string(), tensors, &error)); + GGML_ASSERT(error.find("invalid comfy.gguf.orig_shape") != std::string::npos); + + write_gguf(path, {INT32_MAX, INT32_MAX}); + tensors.clear(); + error.clear(); + GGML_ASSERT(!read_gguf_file(path.string(), tensors, &error)); + GGML_ASSERT(error.find("invalid comfy.gguf.orig_shape") != std::string::npos); + + write_gguf(path, + {256, 96}, + GGML_TYPE_IQ4_XS, + {256, 96}); + tensors.clear(); + error.clear(); + GGML_ASSERT(read_gguf_file(path.string(), tensors, &error)); + GGML_ASSERT(tensors.size() == 1); + GGML_ASSERT(tensors[0].has_comfy_original_shape); + GGML_ASSERT(tensors[0].expected_type == GGML_TYPE_Q4_0); + GGML_ASSERT(tensors[0].ne[0] == 96); + GGML_ASSERT(tensors[0].ne[1] == 256); + + // An unsupported or misaligned source must remain on the explicit F32 fallback + // path and emit the named loader warning rather than being silently remapped. + write_gguf(path, + {256, 100}, + GGML_TYPE_IQ4_XS, + {256, 100}); + tensors.clear(); + error.clear(); + GGML_ASSERT(read_gguf_file(path.string(), tensors, &error)); + GGML_ASSERT(tensors.size() == 1); + GGML_ASSERT(tensors[0].has_comfy_original_shape); + GGML_ASSERT(tensors[0].expected_type == GGML_TYPE_COUNT); + GGML_ASSERT(tensors[0].ne[0] == 100); + GGML_ASSERT(tensors[0].ne[1] == 256); + + const std::string adaln_name = "diffusion_model.blocks.0.adaln_proj.linear.weight"; + int64_t adaln_ne[2] = {32, 32}; + String2TensorStorage mixed_tensors; + mixed_tensors.emplace(adaln_name, + TensorStorage(adaln_name, GGML_TYPE_F32, adaln_ne, 2, 0)); + int64_t unrelated_ne[2] = {32, 32}; + TensorStorage unrelated("vae.encoder.weight", GGML_TYPE_F32, unrelated_ne, 2, 0); + unrelated.has_comfy_original_shape = true; + mixed_tensors.emplace(unrelated.name, unrelated); + + auto config = MiniMaxH3::Config::detect_from_weights(mixed_tensors, "diffusion_model"); + GGML_ASSERT(!config.is_comfyui_layout); + auto adaln = mixed_tensors.find(adaln_name); + GGML_ASSERT(adaln != mixed_tensors.end()); + adaln->second.has_comfy_original_shape = true; + config = MiniMaxH3::Config::detect_from_weights(mixed_tensors, "diffusion_model"); + GGML_ASSERT(config.is_comfyui_layout); + + std::filesystem::remove(path); + return 0; +}