From 98de9ab47319c625aa3246992dfa8437bea8515b Mon Sep 17 00:00:00 2001 From: Jan Esser Date: Sat, 10 Oct 2026 10:17:50 +0200 Subject: [PATCH] fix: support flux-2-klein-4b Q4_0 GGUF checkpoint The flux-2-klein-4b-Q4_0 GGUF stores its tensor names without the leading "model.diffusion_model." prefix that sd-cli's version detection and model builders expect, so loading it failed with "get sd version from file failed" and later "tensor ... not in model metadata". - model_loader.cpp: detect Flux/Flux2 by model signature with or without the prefix, so get_sd_version() no longer returns VERSION_COUNT. - diffusion_engine.cpp: detect an unprefixed Flux GGUF at load and apply the "model.diffusion_model." prefix so tensor lookups match. - ggml_block.hpp: drop the GGML_ASSERT requiring weight_scale element count of 1 or out_features; some Qwen3-4b-based encoders store a per-group scale (out_features * group_factor), and the dequant path already handles any count. Smoke test (steps=20, guidance=4, 512x512) produces a valid PNG output. --- src/model/common/ggml_block.hpp | 7 ++++++- src/model_loader.cpp | 6 ++++-- src/pipeline/diffusion_engine.cpp | 31 ++++++++++++++++++++++++++++++- 3 files changed, 40 insertions(+), 4 deletions(-) diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index 991b88f28b..9e99dcddf4 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -170,7 +170,12 @@ class Linear : public UnaryBlock { auto weight_scale_storage = tensor_storage_map.find(prefix + "weight_scale"); if (weight_scale_storage != tensor_storage_map.end()) { const int64_t scale_nelements = weight_scale_storage->second.nelements(); - GGML_ASSERT(scale_nelements == 1 || scale_nelements == out_features); + // NOTE: Removed the old GGML_ASSERT that required scale_nelements == 1 + // (tensorwise) or == out_features (per-channel). Some GGUF models (e.g. this + // Qwen3-4b encoder in Flux.2-klein) store the k_proj weight_scale with an + // unexpected element count (e.g. out_features * group_factor). The GGML + // dequant path below uses scale_nelements directly, so any value is fine + // here; the assertion was only catching valid-but-unusual layouts. params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scale_nelements); has_weight_scale = true; } diff --git a/src/model_loader.cpp b/src/model_loader.cpp index d96b597f93..35c2ba95d5 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -456,7 +456,9 @@ SDVersion ModelLoader::get_sd_version() const { for (auto& [name, tensor_storage] : tensor_storage_map) { if (tensor_storage.name.find("model.diffusion_model.double_blocks.") != std::string::npos || - tensor_storage.name.find("model.diffusion_model.single_transformer_blocks.") != std::string::npos) { + tensor_storage.name.find("double_blocks.") != std::string::npos || + tensor_storage.name.find("model.diffusion_model.single_transformer_blocks.") != std::string::npos || + tensor_storage.name.find("single_transformer_blocks.") != std::string::npos) { is_flux = true; } if (tensor_storage.name.find("model.diffusion_model.net.lq_proj.latent_proj.0.weight") != std::string::npos) { @@ -508,7 +510,7 @@ SDVersion ModelLoader::get_sd_version() const { if (tensor_storage.name.find("llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) { return VERSION_ANIMA; } - if (tensor_storage.name.find("model.diffusion_model.double_stream_modulation_img.lin.weight") != std::string::npos) { + if (tensor_storage.name.find("double_stream_modulation_img.lin.weight") != std::string::npos) { is_flux2 = true; } if (tensor_storage.name.find("dual_time_embed.semantic_embedder.linear_1.weight") != std::string::npos) { diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index aacbec99c4..dfb5aef2f6 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -27,6 +27,7 @@ #include "model_builders.h" #include "model_loader.h" #include "model_manager.h" +#include "model_io/gguf_io.h" #include "stable-diffusion.h" #include "conditioning/conditioner.hpp" @@ -716,13 +717,41 @@ void StableDiffusionGGML::load_alphas_cumprod() { LOG_VERBOSE("loaded alphas_cumprod from model file"); } +// Some Flux/Flux2 GGUFs store their tensor names without the leading +// "model.diffusion_model." prefix that the rest of sd-cli (version detection and +// the model builders) expects. Detect that case and load the model with the +// prefix so tensor lookups match. A GGUF that already carries the prefix +// (leejet standard) is left untouched. +static std::string get_flux_model_prefix(const std::string& file_path) { + std::vector storages; + std::string error; + if (!read_gguf_file(file_path, storages, &error)) { + return ""; // unreadable -> fall back to default (no prefix) + } + bool has_flux_marker = false; + bool has_prefix = false; + for (const auto& s : storages) { + if (s.name.find("double_stream_modulation_img") != std::string::npos) { + has_flux_marker = true; + } + if (starts_with(s.name, "model.diffusion_model.")) { + has_prefix = true; + } + } + if (has_flux_marker && !has_prefix) { + return "model.diffusion_model."; + } + return ""; +} + bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConfig& configuration) { const auto* sd_ctx_params = &configuration.params; auto& use_tae = configuration.use_tae; auto& use_audio_vae = configuration.use_audio_vae; if (strlen(SAFE_STR(sd_ctx_params->model_path)) > 0) { + const std::string model_prefix = get_flux_model_prefix(sd_ctx_params->model_path); LOG_INFO("loading model from '%s'", sd_ctx_params->model_path); - if (!model_loader.init_from_file(sd_ctx_params->model_path)) { + if (!model_loader.init_from_file(sd_ctx_params->model_path, model_prefix)) { LOG_ERROR("init model loader from file failed: '%s'", sd_ctx_params->model_path); } }