From 936c0cc7fe8debe80bd046974022cc3e2e297a59 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Fri, 18 Sep 2026 13:51:06 +0800 Subject: [PATCH] Align Qwen-Image-2.1 VAE normalization with Diffusers (#7768) Signed-off-by: Tianyu Guo --- .../qwen_image_21/autoencoder_kl_qwenimage21.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/vllm_omni/diffusion/models/qwen_image_21/autoencoder_kl_qwenimage21.py b/vllm_omni/diffusion/models/qwen_image_21/autoencoder_kl_qwenimage21.py index cbaf0aec8..f073b5405 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/autoencoder_kl_qwenimage21.py +++ b/vllm_omni/diffusion/models/qwen_image_21/autoencoder_kl_qwenimage21.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright 2026 The Qwen Team and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -226,7 +229,13 @@ class QwenImage21RMS_norm(nn.Module): self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 def forward(self, x): - return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias + needs_fp32_normalize = x.dtype in (torch.float16, torch.bfloat16) or any( + t in str(x.dtype) for t in ("float4_", "float8_") + ) + normalized = F.normalize(x.float() if needs_fp32_normalize else x, dim=(1 if self.channel_first else -1)).to( + x.dtype + ) + return normalized * self.scale * self.gamma + self.bias class QwenImage21Upsample(nn.Upsample):