From 83bc7766122b439d6e1564ec999e61eac6ab46a7 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 3 Jun 2026 08:32:28 +0800 Subject: [PATCH] [diffusion] optimize: preserve dtype in wanvae nearest upsample (#27077) --- .../runtime/models/vaes/parallel/wan_common_utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/parallel/wan_common_utils.py b/python/sglang/multimodal_gen/runtime/models/vaes/parallel/wan_common_utils.py index 4d4fbde43..2fdb933b5 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/parallel/wan_common_utils.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/parallel/wan_common_utils.py @@ -211,6 +211,8 @@ class WanUpsample(nn.Upsample): """ def forward(self, x): + if current_platform.is_amp_supported(): + return super().forward(x) return super().forward(x.float()).type_as(x)