diff --git a/assets/qwen/qwen-image-2.1-alpha-out2.png b/assets/qwen/qwen-image-2.1-alpha-out2.png index b3db6c892c..8bef08a213 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/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index c4c88a635b..c329a34271 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -465,7 +465,11 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, int d1, int d2, bool force_prec_f32, - bool direct) { + bool direct, + float scale) { + if (scale != 1.f) { + x = ggml_ext_scale(ctx, x, scale); + } if (direct) { int64_t OC = w->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) { b = ggml_reshape_4d(ctx, b, 1, 1, 1, b->ne[0]); // [OC, 1, 1, 1] x = ggml_add_inplace(ctx, x, b); diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index be88c58a3f..5747d625a6 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -154,7 +154,8 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, int d1 = 1, int d2 = 1, bool force_prec_f32 = false, - bool direct = false); + bool direct = false, + float scale = 1.f); // w: [OC,IC, KD, 1 * 1] // x: [N, IC, ID, IH*IW] diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index 73cd09537a..58cff816b5 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -24,6 +24,7 @@ namespace WAN { std::tuple padding; std::tuple dilation; bool bias; + float scale = 1.f; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { auto weight = tensor_storage_map.find(prefix + "weight"); @@ -60,6 +61,10 @@ namespace WAN { dilation(std::move(dilation)), bias(bias) {} + void set_scale(float scale_value) { + scale = scale_value; + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* cache_x = nullptr) { // x: [N*IC, 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, std::get<2>(stride), std::get<1>(stride), 0, 0, 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_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), 0, 0, 0, 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 { blocks["conv2"] = std::shared_ptr(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 all_blocks; + get_all_blocks(all_blocks); + for (auto block : all_blocks) { + if (auto conv = dynamic_cast(block)) { + conv->set_scale(conv_scale); + } else if (auto conv = dynamic_cast(block)) { + conv->set_scale(conv_scale); + } + } + } } static ggml_tensor* patchify(ggml_context* ctx,