diff --git a/assets/qwen/qwen-image-2.1-alpha-out2.png b/assets/qwen/qwen-image-2.1-alpha-out2.png index b3db6c892c..89e32c1e61 100644 Binary files a/assets/qwen/qwen-image-2.1-alpha-out2.png and b/assets/qwen/qwen-image-2.1-alpha-out2.png differ diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 24dc937b58..469fc5afa0 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -460,6 +460,11 @@ ArgOptions SDContextParams::get_options() { "VAE latent format override: auto, flux, sd3, flux2, or wan (default: auto)", 0, &vae_format}, + {"", + "--vae-dtype", + "VAE/TAE inference precision, without this will use weight dtype (f16, f32 or bf16)", + 0, + &vae_dtype}, {"", "--audio-vae", "path to standalone LTX audio vae model", @@ -921,6 +926,7 @@ std::string SDContextParams::to_string() const { << " embeddings_connectors_path: \"" << embeddings_connectors_path << "\",\n" << " vae_path: \"" << vae_path << "\",\n" << " vae_format: \"" << vae_format << "\",\n" + << " vae_dtype: \"" << vae_dtype << "\",\n" << " audio_vae_path: \"" << audio_vae_path << "\",\n" << " audio_encoder_path: \"" << audio_encoder_path << "\",\n" << " taesd_path: \"" << taesd_path << "\",\n" @@ -1018,6 +1024,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { sd_ctx_params.vae_conv_direct = vae_conv_direct; sd_ctx_params.force_sdxl_vae_conv_scale = force_sdxl_vae_conv_scale; sd_ctx_params.vae_format = str_to_vae_format(vae_format); + sd_ctx_params.vae_dtype = vae_dtype.c_str(); sd_ctx_params.max_vram = max_vram.c_str(); sd_ctx_params.disable_prefetch = disable_prefetch; sd_ctx_params.disable_segmented_compute = disable_segmented_compute; diff --git a/examples/common/common.h b/examples/common/common.h index d0d39b738d..2848e4be1f 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -144,6 +144,7 @@ struct SDContextParams { std::string pulid_weights_path; sd_type_t wtype = SD_TYPE_COUNT; std::string tensor_type_rules; + std::string vae_dtype; // VAE/TAE inference precision: f16/f32/bf16; empty uses weight dtype std::string lora_model_dir = "."; std::string hires_upscalers_dir; diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index a2bfdb680c..209d964961 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -248,6 +248,7 @@ typedef struct { const char* tokenizer; // tokenizer.json path or main=FILE,clip-l=FILE,clip-g=FILE assignments; required for PiD and Lens bool sage_attn; int conditioning_cache_size; // Maximum cached conditioning entries per context; 0 disables caching (default: 4) + const char* vae_dtype; // VAE/TAE inference precision: "f16", "f32" or "bf16"; null or empty uses weight dtype } sd_ctx_params_t; typedef struct { diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 74324169ee..6597b9fda7 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -434,7 +434,17 @@ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx, p1 = 0; } - if (direct) { + if (w->type == GGML_TYPE_F32) { + // ggml_conv_2d hardcodes the im2col dtype to f16 for non-bf16 weights, + // which would re-introduce f16 overflow in high-magnitude VAE decoder + // activations. Build the same im2col + mul_mat with explicit f32. + ggml_tensor* im2col = ggml_im2col(ctx, w, x, s0, s1, p0, p1, d0, d1, true, GGML_TYPE_F32); + x = ggml_mul_mat(ctx, + ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]), + ggml_reshape_2d(ctx, w, w->ne[0] * w->ne[1] * w->ne[2], w->ne[3])); + x = ggml_reshape_4d(ctx, x, im2col->ne[1], im2col->ne[2], im2col->ne[3], w->ne[3]); + x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 1, 3, 2)); + } else if (direct) { x = ggml_conv_2d_direct(ctx, w, x, s0, s1, p0, p1, d0, d1); } else { x = ggml_conv_2d(ctx, w, x, s0, s1, p0, p1, d0, d1); @@ -466,11 +476,10 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, int d2, bool force_prec_f32, bool direct) { - if (direct) { - int64_t OC = w->ne[3] / IC; - int64_t N = x->ne[3] / IC; - x = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC); - } else if (force_prec_f32) { + if (force_prec_f32 || w->type == GGML_TYPE_F32) { + // The f32-weight case must bypass ggml_conv_3d: its im2col dtype is + // hardcoded to f16 for non-bf16 weights, which saturates high-magnitude + // VAE decoder activations to inf. ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type); int64_t OC = w->ne[3] / IC; @@ -484,6 +493,10 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, x = ggml_reshape_4d(ctx, x, im2col->ne[1] * im2col->ne[2], OD, N, OC); x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 1, 3, 2)); x = ggml_reshape_4d(ctx, x, im2col->ne[1], im2col->ne[2], OD, OC * N); + } else if (direct) { + int64_t OC = w->ne[3] / IC; + int64_t N = x->ne[3] / IC; + x = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC); } else { // ggml_conv_3d decomposes into GGML_OP_IM2COL_3D, which some backends // (e.g. Metal, see #850) do not implement. Fall back to diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index 59573210b3..29fa150f3b 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -527,6 +527,7 @@ GGMLRunnerContext GGMLRunner::get_context() { runner_ctx.conv3d_direct_enabled = conv3d_direct_enabled; runner_ctx.circular_x_enabled = circular_x_enabled; runner_ctx.circular_y_enabled = circular_y_enabled; + runner_ctx.vae_compute_type = vae_compute_type; runner_ctx.weight_adapter = weight_adapter; runner_ctx.debug_tensors = &debug_tensors; runner_ctx.get_cache_tensor = [this](const std::string& name) { diff --git a/src/core/ggml_runner.h b/src/core/ggml_runner.h index 9084db1097..092c73717e 100644 --- a/src/core/ggml_runner.h +++ b/src/core/ggml_runner.h @@ -77,6 +77,7 @@ struct GGMLRunnerContext { bool circular_y_enabled = false; ggml_tensor* ip_context = nullptr; float ip_scale = 1.0f; + ggml_type vae_compute_type = GGML_TYPE_COUNT; // VAE/TAE conv compute precision; COUNT keeps the stored dtype std::shared_ptr weight_adapter = nullptr; std::vector>* debug_tensors = nullptr; std::function get_cache_tensor; @@ -185,6 +186,7 @@ struct GGMLRunner { bool conv3d_direct_enabled = false; bool circular_x_enabled = false; bool circular_y_enabled = false; + ggml_type vae_compute_type = GGML_TYPE_COUNT; sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_; std::unordered_set params_tensor_set_; @@ -355,6 +357,10 @@ struct GGMLRunner { this->attn_scale = attn_scale; } + void set_vae_compute_type(ggml_type type) { + vae_compute_type = type; + } + void set_conv2d_direct_enabled(bool enabled) { conv2d_direct_enabled = enabled; } diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index 991b88f28b..8dc7003f34 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -427,6 +427,63 @@ class Conv1d : public UnaryBlock { } }; +// VAE/TAE conv weights keep their stored dtype so compute precision can be +// decided per run (see --vae-dtype); other convs (UNet, ControlNet, patch +// embeds, upscalers) stay on the historical hardcoded f16. +static inline bool sd_conv_prefix_is_vae(const std::string& prefix) { + static const std::vector kVaePrefixes = {"first_stage_model", "vae", "tae", "decoder"}; + for (const auto& v : kVaePrefixes) { + if (prefix.size() >= v.size() && prefix.compare(0, v.size(), v) == 0 && + (prefix.size() == v.size() || prefix[v.size()] == '.')) { + return true; + } + } + return false; +} + +static inline bool sd_conv_type_computable(ggml_type t) { + return t == GGML_TYPE_F16 || t == GGML_TYPE_F32 || t == GGML_TYPE_BF16; +} + +static inline ggml_type sd_vae_conv_param_type(const String2TensorStorage& map, const std::string& key) { + auto it = map.find(key); + if (it == map.end()) { + return GGML_TYPE_F16; + } + const TensorStorage& ts = it->second; + ggml_type t = ts.expected_type != GGML_TYPE_COUNT ? ts.expected_type : ts.type; + if (sd_conv_type_computable(t)) { + return t; + } + // Block-quantized types need ne[0] % blck_size == 0, but conv weights have + // ne[0] = kernel width, so they cannot stay quantized in a conv param and + // keep the historical load-time f16 expansion. int8-tensorwise needs its + // weight_scale applied by dedicated kernels that conv paths do not have. + if (ggml_is_quantized(t) && !ts.is_int8_tensorwise && ts.ne[0] % ggml_blck_size(t) == 0) { + return t; + } + return GGML_TYPE_F16; +} + +static inline ggml_type sd_conv_target_type(GGMLRunnerContext* ctx, ggml_type wtype) { + if (ctx->vae_compute_type != GGML_TYPE_COUNT) { + return ctx->vae_compute_type; + } + return sd_conv_type_computable(wtype) ? wtype : GGML_TYPE_F16; +} + +// Quantized sources route through f32: backends only register direct +// CPY kernels for quantized -> f32, not quantized -> f16/bf16. +static inline ggml_tensor* sd_conv_cast_weight(ggml_context* ctx, ggml_tensor* w, ggml_type target) { + if (w->type == target) { + return w; + } + if (ggml_is_quantized(w->type) && target != GGML_TYPE_F32) { + w = ggml_cast(ctx, w, GGML_TYPE_F32); + } + return ggml_cast(ctx, w, target); +} + class Conv2d : public UnaryBlock { protected: int64_t in_channels; @@ -440,9 +497,10 @@ class Conv2d : public UnaryBlock { std::string prefix; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override { - this->prefix = prefix; - enum ggml_type wtype = GGML_TYPE_F16; - params["weight"] = ggml_new_tensor_4d(ctx, wtype, kernel_size.second, kernel_size.first, in_channels, out_channels); + this->prefix = prefix; + enum ggml_type wtype = + sd_conv_prefix_is_vae(prefix) ? sd_vae_conv_param_type(tensor_storage_map, prefix + "weight") : GGML_TYPE_F16; + params["weight"] = ggml_new_tensor_4d(ctx, wtype, kernel_size.second, kernel_size.first, in_channels, out_channels); if (bias) { enum ggml_type wtype = GGML_TYPE_F32; params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_channels); @@ -475,6 +533,7 @@ class Conv2d : public UnaryBlock { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { ggml_tensor* w = params["weight"]; + w = sd_conv_cast_weight(ctx->ggml_ctx, w, sd_conv_target_type(ctx, w->type)); ggml_tensor* b = nullptr; if (bias) { b = params["bias"]; @@ -525,9 +584,10 @@ class Conv2d_grouped : public UnaryBlock { std::string prefix; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override { - this->prefix = prefix; - enum ggml_type wtype = GGML_TYPE_F16; - params["weight"] = ggml_new_tensor_4d(ctx, wtype, kernel_size.second, kernel_size.first, in_channels / groups, out_channels); + this->prefix = prefix; + enum ggml_type wtype = + sd_conv_prefix_is_vae(prefix) ? sd_vae_conv_param_type(tensor_storage_map, prefix + "weight") : GGML_TYPE_F16; + params["weight"] = ggml_new_tensor_4d(ctx, wtype, kernel_size.second, kernel_size.first, in_channels / groups, out_channels); if (bias) { enum ggml_type wtype = GGML_TYPE_F32; params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_channels); @@ -562,6 +622,7 @@ class Conv2d_grouped : public UnaryBlock { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { ggml_tensor* w = params["weight"]; + w = sd_conv_cast_weight(ctx->ggml_ctx, w, sd_conv_target_type(ctx, w->type)); ggml_tensor* b = nullptr; if (bias) { b = params["bias"]; @@ -682,14 +743,15 @@ class Conv3d : public UnaryBlock { std::string prefix; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override { - this->prefix = prefix; - enum ggml_type wtype = GGML_TYPE_F16; - params["weight"] = ggml_new_tensor_4d(ctx, - wtype, - std::get<2>(kernel_size), - std::get<1>(kernel_size), - std::get<0>(kernel_size), - in_channels * out_channels); + this->prefix = prefix; + enum ggml_type wtype = + sd_conv_prefix_is_vae(prefix) ? sd_vae_conv_param_type(tensor_storage_map, prefix + "weight") : GGML_TYPE_F16; + params["weight"] = ggml_new_tensor_4d(ctx, + wtype, + std::get<2>(kernel_size), + std::get<1>(kernel_size), + std::get<0>(kernel_size), + in_channels * out_channels); if (bias) { params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); } @@ -715,6 +777,7 @@ class Conv3d : public UnaryBlock { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { ggml_tensor* w = params["weight"]; + w = sd_conv_cast_weight(ctx->ggml_ctx, w, sd_conv_target_type(ctx, w->type)); ggml_tensor* b = nullptr; if (ctx->weight_adapter) { w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index da59cfab61..af232426e4 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -277,6 +277,11 @@ namespace LLM { if (!starts_with(name, prefix)) { continue; } + // Quant scale tensors mirror the weight name and can share its suffix with a + // different shape; never use them for config detection. + if (contains(name, "weight_scale")) { + continue; + } size_t pos = name.find("visual."); if (pos != std::string::npos) { config.have_vision_weight = true; @@ -332,7 +337,9 @@ namespace LLM { } } } - if (contains(name, "embed_tokens.weight")) { + // ends_with: int8 checkpoints also carry embed_tokens.weight_scale tensors whose + // shape would otherwise be mistaken for the embedding matrix. + if (ends_with(name, "embed_tokens.weight")) { config.hidden_size = tensor_storage.ne[0]; config.vocab_size = tensor_storage.ne[1]; } diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index 8d14cf9570..2fc461f245 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -33,8 +33,13 @@ namespace WAN { std::get<0>(kernel_size) = 1; std::get<0>(padding) = 0; } + // VAE conv weights keep their stored dtype so compute precision can + // be decided per run; the decoder residual stream can exceed the + // f16 range, so forcing f16 here saturates im2col output to inf. + ggml_type wtype = sd_conv_prefix_is_vae(prefix) ? sd_vae_conv_param_type(tensor_storage_map, prefix + "weight") + : GGML_TYPE_F16; params["weight"] = ggml_new_tensor_4d(ctx, - GGML_TYPE_F16, + wtype, std::get<2>(kernel_size), std::get<1>(kernel_size), std::get<0>(kernel_size), @@ -64,6 +69,7 @@ namespace WAN { // x: [N*IC, ID, IH, IW] // result: x: [N*OC, ID, IH, IW] ggml_tensor* w = params["weight"]; + w = sd_conv_cast_weight(ctx->ggml_ctx, w, sd_conv_target_type(ctx, w->type)); ggml_tensor* b = nullptr; if (bias) { b = params["bias"]; diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index 6f1b598e88..a9c45b7b87 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -1568,6 +1568,12 @@ std::string convert_tensor_name(std::string name, SDVersion version) { replace_with_prefix_map(name, prefix_map); if (version == VERSION_QWEN_IMAGE_2_1 || sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) { + // Recent transformers Qwen3-VL checkpoints nest the text backbone under + // model.language_model.* instead of model.*. + const std::string hf_lm_prefix = "text_encoders.llm.model.language_model."; + if (starts_with(name, hf_lm_prefix)) { + name = "text_encoders.llm.model." + name.substr(hf_lm_prefix.size()); + } const std::string hf_vision_prefix = "text_encoders.llm.model.visual."; if (starts_with(name, hf_vision_prefix)) { name = "text_encoders.llm.visual." + name.substr(hf_vision_prefix.size()); diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index 9b8819a05f..0dd1138d4c 100644 --- a/src/pipeline/model_builders.cpp +++ b/src/pipeline/model_builders.cpp @@ -615,6 +615,26 @@ namespace sd::model_builders { result.preview->set_conv3d_direct_enabled(true); } } + { + ggml_type vae_compute_type = GGML_TYPE_COUNT; + if (SAFE_STR(sd_ctx_params->vae_dtype)[0] != '\0') { + const std::string dtype = sd_ctx_params->vae_dtype; + if (dtype == "f16" || dtype == "f32" || dtype == "bf16") { + vae_compute_type = dtype == "f16" ? GGML_TYPE_F16 + : dtype == "f32" ? GGML_TYPE_F32 + : GGML_TYPE_BF16; + } else { + LOG_WARN("invalid --vae-dtype '%s' (expected f16, f32 or bf16), using weight dtype", dtype.c_str()); + } + } + if (vae_compute_type != GGML_TYPE_COUNT) { + LOG_INFO("VAE/TAE compute dtype: %s", ggml_type_name(vae_compute_type)); + result.vae->set_vae_compute_type(vae_compute_type); + if (result.preview) { + result.preview->set_vae_compute_type(vae_compute_type); + } + } + } if (result.vae) { result.vae->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale); } diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index b13c055f3d..7d7d46e817 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -398,7 +398,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "sage_attn: %s\n" "linear_scale: %g\n" "attn_scale: %g\n" - "vae_format: %s\n", + "vae_format: %s\n" + "vae_dtype: %s\n", SAFE_STR(sd_ctx_params->model_path), SAFE_STR(sd_ctx_params->clip_l_path), SAFE_STR(sd_ctx_params->clip_g_path), @@ -439,7 +440,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { BOOL_STR(sd_ctx_params->sage_attn), sd_ctx_params->linear_scale, sd_ctx_params->attn_scale, - sd_vae_format_name(sd_ctx_params->vae_format)); + sd_vae_format_name(sd_ctx_params->vae_format), + SAFE_STR(sd_ctx_params->vae_dtype)); return buf; }