fix: scale Qwen Image 2.1 VAE convolutions (#2054)

This commit is contained in:
leejet 2026-09-25 15:52:35 +08:00 committed by GitHub
parent b167b942f7
commit 0a9340c599
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 30 additions and 4 deletions

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.7 MiB

After

Width:  |  Height:  |  Size: 1.1 MiB

View File

@ -465,7 +465,11 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
int d1, int d1,
int d2, int d2,
bool force_prec_f32, bool force_prec_f32,
bool direct) { bool direct,
float scale) {
if (scale != 1.f) {
x = ggml_ext_scale(ctx, x, scale);
}
if (direct) { if (direct) {
int64_t OC = w->ne[3] / IC; int64_t OC = w->ne[3] / IC;
int64_t N = x->ne[3] / IC; int64_t N = x->ne[3] / IC;
@ -502,6 +506,9 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
} }
} }
if (scale != 1.f) {
x = ggml_ext_scale(ctx, x, 1.f / scale);
}
if (b != nullptr) { if (b != nullptr) {
b = ggml_reshape_4d(ctx, b, 1, 1, 1, b->ne[0]); // [OC, 1, 1, 1] b = ggml_reshape_4d(ctx, b, 1, 1, 1, b->ne[0]); // [OC, 1, 1, 1]
x = ggml_add_inplace(ctx, x, b); x = ggml_add_inplace(ctx, x, b);

View File

@ -154,7 +154,8 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
int d1 = 1, int d1 = 1,
int d2 = 1, int d2 = 1,
bool force_prec_f32 = false, bool force_prec_f32 = false,
bool direct = false); bool direct = false,
float scale = 1.f);
// w: [OC,IC, KD, 1 * 1] // w: [OC,IC, KD, 1 * 1]
// x: [N, IC, ID, IH*IW] // x: [N, IC, ID, IH*IW]

View File

@ -24,6 +24,7 @@ namespace WAN {
std::tuple<int, int, int> padding; std::tuple<int, int, int> padding;
std::tuple<int, int, int> dilation; std::tuple<int, int, int> dilation;
bool bias; bool bias;
float scale = 1.f;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
auto weight = tensor_storage_map.find(prefix + "weight"); auto weight = tensor_storage_map.find(prefix + "weight");
@ -60,6 +61,10 @@ namespace WAN {
dilation(std::move(dilation)), dilation(std::move(dilation)),
bias(bias) {} bias(bias) {}
void set_scale(float scale_value) {
scale = scale_value;
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* cache_x = nullptr) { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* cache_x = nullptr) {
// x: [N*IC, ID, IH, IW] // x: [N*IC, ID, IH, IW]
// result: x: [N*OC, ID, IH, IW] // result: x: [N*OC, ID, IH, IW]
@ -93,14 +98,14 @@ namespace WAN {
x2 = ggml_ext_conv_2d(ctx->ggml_ctx, x2, w2, b, x2 = ggml_ext_conv_2d(ctx->ggml_ctx, x2, w2, b,
std::get<2>(stride), std::get<1>(stride), 0, 0, std::get<2>(stride), std::get<1>(stride), 0, 0,
std::get<2>(dilation), std::get<1>(dilation), std::get<2>(dilation), std::get<1>(dilation),
ctx->conv2d_direct_enabled); ctx->conv2d_direct_enabled, false, false, scale);
return ggml_reshape_4d(ctx->ggml_ctx, x2, x2->ne[0], x2->ne[1], 1, out_channels); return ggml_reshape_4d(ctx->ggml_ctx, x2, x2->ne[0], x2->ne[1], 1, out_channels);
} }
return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels, return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
std::get<2>(stride), std::get<1>(stride), std::get<0>(stride), std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
0, 0, 0, 0, 0, 0,
std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation), std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation),
false, ctx->conv3d_direct_enabled); false, ctx->conv3d_direct_enabled, scale);
} }
}; };
@ -1117,6 +1122,19 @@ namespace WAN {
} else { } else {
blocks["conv2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, z_dim, {1, 1, 1})); blocks["conv2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, z_dim, {1, 1, 1}));
} }
if (version == VERSION_QWEN_IMAGE_2_1) {
// Keep large VAE activations within the FP16 convolution range.
const float conv_scale = 1.f / 128.f;
std::vector<GGMLBlock*> all_blocks;
get_all_blocks(all_blocks);
for (auto block : all_blocks) {
if (auto conv = dynamic_cast<Conv2d*>(block)) {
conv->set_scale(conv_scale);
} else if (auto conv = dynamic_cast<CausalConv3d*>(block)) {
conv->set_scale(conv_scale);
}
}
}
} }
static ggml_tensor* patchify(ggml_context* ctx, static ggml_tensor* patchify(ggml_context* ctx,