minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Diffusion] Port the Wan VAE decoder fast paths to the Qwen-Image VAE (#38020)

Browse Source 
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
This commit is contained in:
Dayuxiaoshui
2026-09-05 18:02:33 +08:00
committed by GitHub
co-authored by Xiaoyu Zhang
parent 0948e6ebed
commit 50c1bf0db0
5 changed files with 446 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md
+11 -3
View File
@@ -47,6 +47,7 @@ framework-specific optimization workflow.
- `python/sglang/multimodal_gen/runtime/models/vaes/fast_path_gate.py`
- `python/sglang/multimodal_gen/runtime/models/vaes/flux2_vae_cuda_opt.py`
- `python/sglang/multimodal_gen/runtime/models/vaes/wan_vae_cuda_opt.py`
- `python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py`
- `python/sglang/multimodal_gen/runtime/breakable_cuda_graph/runner.py`
- `test/registered/kernels/ops/diffusion/test_qwen_image_modulation.py`
- `test/registered/kernels/ops/diffusion/test_group_norm_silu.py`
@@ -128,8 +129,10 @@ framework-specific optimization workflow.
Wan cublasLt/NVFP4 GELU, Qwen added-QKV, GLM/Qwen/Hunyuan/LTX fused GELU,
LTX RMSNorm+modulate, Hunyuan QK RMSNorm, Ideogram gated RMSNorm,
LingBot RMSNorm, SANA-Video linear attention, generic KL VAE
decoder rewrites used by FLUX.1/FLUX.2/Z-Image/SD3, and Wan VAE
RMSNorm+SiLU.
decoder rewrites used by FLUX.1/FLUX.2/Z-Image/SD3, and Wan / Qwen-Image
VAE RMSNorm+SiLU (the Qwen-Image VAE is the Wan 2.1 VAE; its gate also
re-expresses the `Resample` upsample input with canonical NHWC strides so
the 2D conv runs channels_last end-to-end).
- Do not confuse request `--quality` with `--output-quality`, which controls
output-file compression rather than model math.
- Validation: `test_quality_gate.py`, `test_fused_ln_modulate.py`,
@@ -223,7 +226,12 @@ framework-specific optimization workflow.
- Numerical contract: these are bit-exact data-movement / same-order-add
replacements and run independently of the request-gated Wan RMSNorm+SiLU
path. Unsupported layouts or padding fall back to the aten chain.
- Validation: `test/registered/kernels/ops/diffusion/test_wan_causal_cache.py`.
- The Qwen-Image VAE (`autoencoder_kl_qwenimage.py`) uses the same
`cat_pad_channels_last_3d` + compact-cache helper for every causal conv
slot; single-frame image decodes keep the compact cache at the reference
size (one frame) so peak memory does not grow.
- Validation: `test/registered/kernels/ops/diffusion/test_wan_causal_cache.py`
and the `test_qwen_vae_*` cases in `test_model_fast_paths.py`.
15. Helios paired transposed RoPE
- Kernel: `fused_inplace_helios_qk_rope`.
python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py
+110 -103
View File
@@ -37,13 +37,90 @@ from sglang.multimodal_gen.runtime.models.vaes.common import (
has_decode_parallel_world,
should_run_spatial_shard_parallel_decode,
)
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
if current_platform.is_cuda():
try:
from sglang.kernels.ops.diffusion import cat_pad_channels_last_3d
except ImportError: # pragma: no cover
cat_pad_channels_last_3d = None
else:
cat_pad_channels_last_3d = None
logger = init_logger(__name__) # pylint: disable=invalid-name
CACHE_T = 2
def _fused_conv_cache_supported(conv: nn.Module, x: torch.Tensor) -> bool:
"""The fused cat+pad writes channels_last_3d, so it only pays off (and is
only exercised) when the conv weights are channels_last_3d too; the kernel
symbol being present already implies a CUDA platform."""
return (
cat_pad_channels_last_3d is not None
and type(conv) is QwenImageCausalConv3d
and x.dim() == 5
and x.is_cuda
and conv.weight.is_contiguous(memory_format=torch.channels_last_3d)
and not torch.compiler.is_compiling()
)
def _run_cached_causal_conv(
conv: nn.Module,
x: torch.Tensor,
cache_list: list,
idx: int,
) -> torch.Tensor:
"""Run one causal conv, consuming and refreshing its feature-cache slot.
Same contract as the Wan VAE helper (this VAE is the Wan 2.1 VAE): the
fast path builds the conv input (cache frames + hidden state + padding)
directly in channels_last_3d with one kernel and takes the next cache
entry as one compact copy of that input's unpadded tail, instead of the
per-chunk ``clone``/``cat``/``F.pad``/relayout bookkeeping. The compact
cache holds exactly the reference cache values, so fused and fallback
chunks can interleave. Pure data movement plus zero fill, so the conv
input is bit-identical to the aten chain. Falls back to the original op
chain whenever the fused kernel does not support the request.
"""
cache = cache_list[idx]
is_rep = isinstance(cache, str) # "Rep" marker from QwenImageResample
payload = None if is_rep else cache
if _fused_conv_cache_supported(conv, x) and (
payload is None or (payload.device == x.device and payload.dtype == x.dtype)
):
# Emit exactly the frames the reference bookkeeping would keep: the
# first chunk of a clip stores only its own tail (``x[:, :, -CACHE_T:]``),
# later chunks and the "Rep" slot always hold ``CACHE_T`` frames. This
# keeps single-frame (image) decodes from pinning a second full-size
# frame per conv site.
keep_t = (
CACHE_T if (payload is not None or is_rep) else min(CACHE_T, x.shape[2])
)
pair = cat_pad_channels_last_3d(x, payload, conv._padding, keep_cache_t=keep_t)
if pair is not None:
inp, cache_list[idx] = pair
return nn.Conv3d.forward(conv, inp)
# Original aten path (bit-identical bookkeeping).
cache_x = x[:, :, -CACHE_T:, :, :].clone()
if cache_x.shape[2] < 2 and payload is not None:
# cache last frame of last two chunk
cache_x = torch.cat(
[payload[:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
dim=2,
)
elif cache_x.shape[2] < 2 and is_rep:
cache_x = torch.cat(
[torch.zeros_like(cache_x).to(cache_x.device), cache_x],
dim=2,
)
out = conv(x) if payload is None else conv(x, payload)
cache_list[idx] = cache_x
return out
class QwenImageCausalConv3d(nn.Conv3d):
r"""
A custom 3D causal convolution layer with feature caching support.
@@ -88,6 +165,18 @@ class QwenImageCausalConv3d(nn.Conv3d):
def forward(self, x, cache_x=None):
padding = list(self._padding)
if (
any(padding)
and _fused_conv_cache_supported(self, x)
and (
cache_x is None
or (cache_x.device == x.device and cache_x.dtype == x.dtype)
)
):
# Bit-exact: cat + pad + channels_last_3d relayout in one pass.
inp = cat_pad_channels_last_3d(x, cache_x, padding)
if inp is not None:
return super().forward(inp)
x = causal_conv3d_cat_pad(x, cache_x, padding)
return super().forward(x)
@@ -138,6 +227,10 @@ class QwenImageUpsample(nn.Upsample):
"""
def forward(self, x):
if current_platform.is_amp_supported():
# nearest-exact is a pure gather, so bf16/fp16 is bit-identical to
# the fp32 round trip; skip its three full-size copies.
return super().forward(x)
return super().forward(x.float()).type_as(x)
@@ -220,36 +313,7 @@ class QwenImageResample(nn.Module):
feat_cache[idx] = "Rep"
feat_idx[0] += 1
else:
cache_x = x[:, :, -CACHE_T:, :, :].clone()
if (
cache_x.shape[2] < 2
and feat_cache[idx] is not None
and feat_cache[idx] != "Rep"
):
# cache last frame of last two chunk
cache_x = torch.cat(
[
feat_cache[idx][:, :, -1, :, :]
.unsqueeze(2)
.to(cache_x.device),
cache_x,
],
dim=2,
)
if (
cache_x.shape[2] < 2
and feat_cache[idx] is not None
and feat_cache[idx] == "Rep"
):
cache_x = torch.cat(
[torch.zeros_like(cache_x).to(cache_x.device), cache_x],
dim=2,
)
if feat_cache[idx] == "Rep":
x = self.time_conv(x)
else:
x = self.time_conv(x, feat_cache[idx])
feat_cache[idx] = cache_x
x = _run_cached_causal_conv(self.time_conv, x, feat_cache, idx)
feat_idx[0] += 1
x = x.reshape(b, 2, c, t, h, w)
@@ -327,18 +391,7 @@ class QwenImageResidualBlock(nn.Module):
if feat_cache is not None:
idx = feat_idx[0]
cache_x = x[:, :, -CACHE_T:, :, :].clone()
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
cache_x = torch.cat(
[
feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device),
cache_x,
],
dim=2,
)
x = self.conv1(x, feat_cache[idx])
feat_cache[idx] = cache_x
x = _run_cached_causal_conv(self.conv1, x, feat_cache, idx)
feat_idx[0] += 1
else:
x = self.conv1(x)
@@ -352,18 +405,7 @@ class QwenImageResidualBlock(nn.Module):
if feat_cache is not None:
idx = feat_idx[0]
cache_x = x[:, :, -CACHE_T:, :, :].clone()
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
cache_x = torch.cat(
[
feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device),
cache_x,
],
dim=2,
)
x = self.conv2(x, feat_cache[idx])
feat_cache[idx] = cache_x
x = _run_cached_causal_conv(self.conv2, x, feat_cache, idx)
feat_idx[0] += 1
else:
x = self.conv2(x)
@@ -421,7 +463,16 @@ class QwenImageAttentionBlock(nn.Module):
x = x.view(batch_size, time, channels, height, width)
x = x.permute(0, 2, 1, 3, 4)
x = x + identity
gate = getattr(self, "_sgl_gate", None)
if gate is not None and gate.enabled and not torch.compiler.is_compiling():
# ``identity`` carries the decoder's channels_last_3d layout;
# putting it first makes the sum inherit that layout so the next
# block's fused RMSNorm+SiLU applies. The add is commutative, but
# the layout changes the reduction order of an eager norm that
# follows under fp32 autocast, so this stays quality-gated.
x = identity + x
else:
x = x + identity
if self.spatial_parallel:
x = chunk_height_for_parallel_decode(x)
return x
@@ -565,18 +616,7 @@ class QwenImageEncoder3d(nn.Module):
def forward(self, x, feat_cache=None, feat_idx=[0]):
if feat_cache is not None:
idx = feat_idx[0]
cache_x = x[:, :, -CACHE_T:, :, :].clone()
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
# cache last frame of last two chunk
cache_x = torch.cat(
[
feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device),
cache_x,
],
dim=2,
)
x = self.conv_in(x, feat_cache[idx])
feat_cache[idx] = cache_x
x = _run_cached_causal_conv(self.conv_in, x, feat_cache, idx)
feat_idx[0] += 1
else:
x = self.conv_in(x)
@@ -596,18 +636,7 @@ class QwenImageEncoder3d(nn.Module):
x = self.nonlinearity(x)
if feat_cache is not None:
idx = feat_idx[0]
cache_x = x[:, :, -CACHE_T:, :, :].clone()
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
# cache last frame of last two chunk
cache_x = torch.cat(
[
feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device),
cache_x,
],
dim=2,
)
x = self.conv_out(x, feat_cache[idx])
feat_cache[idx] = cache_x
x = _run_cached_causal_conv(self.conv_out, x, feat_cache, idx)
feat_idx[0] += 1
else:
x = self.conv_out(x)
@@ -828,18 +857,7 @@ class QwenImageDecoder3d(nn.Module):
## conv1
if feat_cache is not None:
idx = feat_idx[0]
cache_x = x[:, :, -CACHE_T:, :, :].clone()
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
# cache last frame of last two chunk
cache_x = torch.cat(
[
feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device),
cache_x,
],
dim=2,
)
x = self.conv_in(x, feat_cache[idx])
feat_cache[idx] = cache_x
x = _run_cached_causal_conv(self.conv_in, x, feat_cache, idx)
feat_idx[0] += 1
else:
x = self.conv_in(x)
@@ -856,18 +874,7 @@ class QwenImageDecoder3d(nn.Module):
x = self.nonlinearity(x)
if feat_cache is not None:
idx = feat_idx[0]
cache_x = x[:, :, -CACHE_T:, :, :].clone()
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
# cache last frame of last two chunk
cache_x = torch.cat(
[
feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device),
cache_x,
],
dim=2,
)
x = self.conv_out(x, feat_cache[idx])
feat_cache[idx] = cache_x
x = _run_cached_causal_conv(self.conv_out, x, feat_cache, idx)
feat_idx[0] += 1
else:
x = self.conv_out(x)
python/sglang/multimodal_gen/runtime/models/vaes/wan_vae_cuda_opt.py
+166 -19
View File
@@ -1,5 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
"""CUDA fast path for the Wan VAE decoder (AutoencoderKLWan).
"""CUDA fast path for the Wan-family VAE decoders (AutoencoderKLWan and the
Qwen-Image VAE, which is the Wan 2.1 VAE under other class names).
Fuses every decoder ``WanRMS_norm -> SiLU`` chain into one Triton kernel on
the channels_last_3d layout. Wrappers are installed once at VAE load and
@@ -58,6 +59,62 @@ class FusedWanRMSNormSiLU(nn.Module):
return F.silu(F.normalize(x, dim=1) * self.scale * self.gamma + self.bias)
class GatedChannelsLastUpsample(nn.Module):
"""Wan-style ``Resample`` upsample that keeps the decoder channels_last.
``Resample.forward`` reaches its 2D ``Upsample -> Conv2d`` through
``x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)``. On a channels_last_3d
tensor with ``b * t == 1`` that view keeps a degenerate batch stride which
``is_contiguous(channels_last)`` accepts but aten's
``suggest_memory_format`` does not, so the upsample and the conv2d run
their NCHW kernels (cuDNN converts back and forth internally) and every
up block starts with an NCDHW tensor that the fused RMSNorm+SiLU cannot
take. With the gate on, re-express the same memory with canonical NHWC
strides (a pure view) so the conv2d runs channels_last end-to-end. The
NHWC conv2d is not guaranteed to pick the same cuDNN algorithm as the NCHW
one, so the stride canonicalisation stays behind the quality gate; with
the gate off the layout is left exactly as the eager module sees it.
"""
def __init__(self, upsample: nn.Upsample, gate: VaeFastPathGate) -> None:
super().__init__()
self._sgl_upsample = upsample
self._sgl_gate = gate
def forward(self, x: torch.Tensor) -> torch.Tensor:
up = self._sgl_upsample
if torch.compiler.is_compiling() or x.dim() != 4:
return up(x)
_, c, h, w = x.shape
canonical = (h * w * c, 1, w * c, c)
# ``is_contiguous(channels_last)`` guarantees every dim of size > 1
# already has its canonical stride; only size-1 dims (the merged
# batch dim here) can carry a stray one, and their stride is never
# used to address memory. So ``as_strided`` to the canonical strides
# is a pure re-labelling of the same elements, for any batch size.
if (
self._sgl_gate.enabled
and x.stride() != canonical
and x.is_contiguous(memory_format=torch.channels_last)
):
x = x.as_strided(x.shape, canonical)
return up(x)
def _install_module_gates(
decoder: nn.Module, gate: VaeFastPathGate, module_classes: tuple[type, ...]
) -> int:
"""Hand the gate to layout-sensitive modules that consult ``_sgl_gate`` in
their own forward (the attention block's residual operand order, so the
sum keeps the channels_last_3d layout)."""
count = 0
for m in decoder.modules():
if type(m) in module_classes:
m._sgl_gate = gate
count += 1
return count
def _is_plain_silu(act: object) -> bool:
return isinstance(act, nn.SiLU) and not act.inplace
@@ -70,34 +127,42 @@ def _norm_fusable(norm: object, wan_rms_norm_cls: type) -> bool:
)
def _install_norm_silu(decoder: nn.Module, gate: VaeFastPathGate) -> int | None:
"""Wrap every decoder ``WanRMS_norm -> SiLU`` chain; ``None`` = fail closed."""
from sglang.multimodal_gen.runtime.models.vaes.wanvae import (
WanResidualBlock,
WanRMS_norm,
)
def _install_norm_silu(
decoder: nn.Module,
gate: VaeFastPathGate,
*,
residual_block_cls: type,
rms_norm_cls: type,
label: str,
) -> int | None:
"""Wrap every decoder ``RMS_norm -> SiLU`` chain; ``None`` = fail closed.
res_blocks = [m for m in decoder.modules() if isinstance(m, WanResidualBlock)]
``residual_block_cls`` / ``rms_norm_cls`` name the model's own Wan-style
residual block and channel-first RMSNorm (``WanResidualBlock`` /
``WanRMS_norm`` or their Qwen-Image twins).
"""
res_blocks = [m for m in decoder.modules() if isinstance(m, residual_block_cls)]
eligible = [
m
for m in res_blocks
if type(m) is WanResidualBlock
if type(m) is residual_block_cls
and _is_plain_silu(m.nonlinearity)
and _norm_fusable(m.norm1, WanRMS_norm)
and _norm_fusable(m.norm2, WanRMS_norm)
and _norm_fusable(m.norm1, rms_norm_cls)
and _norm_fusable(m.norm2, rms_norm_cls)
]
if len(eligible) != len(res_blocks):
logger.warning(
"Wan VAE: %d/%d residual blocks non-standard; skipping fast path.",
"%s: %d/%d residual blocks non-standard; skipping fast path.",
label,
len(res_blocks) - len(eligible),
len(res_blocks),
)
return None
if not (
_norm_fusable(getattr(decoder, "norm_out", None), WanRMS_norm)
_norm_fusable(getattr(decoder, "norm_out", None), rms_norm_cls)
and _is_plain_silu(getattr(decoder, "nonlinearity", None))
):
logger.warning("Wan VAE: non-standard output head; skipping fast path.")
logger.warning("%s: non-standard output head; skipping fast path.", label)
return None
count = 0
@@ -111,32 +176,114 @@ def _install_norm_silu(decoder: nn.Module, gate: VaeFastPathGate) -> int | None:
return count + 1
def _install_channels_last_upsample(
decoder: nn.Module, gate: VaeFastPathGate, upsample_cls: type
) -> int:
"""Wrap the 2D nearest upsample of every ``Resample`` (index 0 of its
``resample`` Sequential, no parameters, so state_dict names are kept)."""
count = 0
for m in decoder.modules():
seq = getattr(m, "resample", None)
if (
isinstance(seq, nn.Sequential)
and len(seq) >= 1
and type(seq[0]) is upsample_cls
):
seq[0] = GatedChannelsLastUpsample(seq[0], gate)
count += 1
return count
def maybe_optimize_wan_vae(vae: nn.Module) -> nn.Module:
"""Install the quality-gated CUDA Wan VAE decoder fast path."""
from sglang.multimodal_gen.runtime.models.vaes.wanvae import (
AutoencoderKLWan,
WanDecoder3d,
WanResidualBlock,
WanRMS_norm,
)
if not isinstance(vae, AutoencoderKLWan):
return vae
return _maybe_optimize_wan_family_vae(
vae,
decoder_cls=WanDecoder3d,
residual_block_cls=WanResidualBlock,
rms_norm_cls=WanRMS_norm,
label="Wan VAE",
)
def maybe_optimize_qwen_image_vae(vae: nn.Module) -> nn.Module:
"""Install the quality-gated CUDA Qwen-Image VAE decoder fast path.
The Qwen-Image VAE is the Wan 2.1 VAE (same ``F.normalize``-based
channel-first RMSNorm, same residual block layout), so it takes the same
``RMSNorm -> SiLU`` fusion under the same gate.
"""
from sglang.multimodal_gen.runtime.models.vaes.autoencoder_kl_qwenimage import (
AutoencoderKLQwenImage,
QwenImageAttentionBlock,
QwenImageDecoder3d,
QwenImageResidualBlock,
QwenImageRMS_norm,
QwenImageUpsample,
)
if not isinstance(vae, AutoencoderKLQwenImage):
return vae
return _maybe_optimize_wan_family_vae(
vae,
decoder_cls=QwenImageDecoder3d,
residual_block_cls=QwenImageResidualBlock,
rms_norm_cls=QwenImageRMS_norm,
upsample_cls=QwenImageUpsample,
gated_module_classes=(QwenImageAttentionBlock,),
label="Qwen-Image VAE",
)
def _maybe_optimize_wan_family_vae(
vae: nn.Module,
*,
decoder_cls: type,
residual_block_cls: type,
rms_norm_cls: type,
label: str,
upsample_cls: type | None = None,
gated_module_classes: tuple[type, ...] = (),
) -> nn.Module:
decoder = getattr(vae, "decoder", None)
if type(decoder) is not WanDecoder3d:
if type(decoder) is not decoder_cls:
return vae
if decoder.use_parallel_decode and decoder.world_size > 1:
logger.info("Wan VAE: spatial-parallel decode; skipping fast path.")
logger.info("%s: spatial-parallel decode; skipping fast path.", label)
return vae
if not _HAS_TRITON:
logger.warning("Wan VAE: Triton unavailable; skipping fast path.")
logger.warning("%s: Triton unavailable; skipping fast path.", label)
return vae
gate = VaeFastPathGate()
n_norm = _install_norm_silu(decoder, gate)
n_norm = _install_norm_silu(
decoder,
gate,
residual_block_cls=residual_block_cls,
rms_norm_cls=rms_norm_cls,
label=label,
)
if n_norm is None:
return vae
n_up = 0
if upsample_cls is not None:
n_up = _install_channels_last_upsample(decoder, gate, upsample_cls)
if gated_module_classes:
_install_module_gates(decoder, gate, gated_module_classes)
register_vae_fast_path_gate(vae, gate)
logger.info(
"Wan VAE: installed quality-gated fast path (%d RMSNorm+SiLU fusions).",
"%s: installed quality-gated fast path (%d RMSNorm+SiLU fusions, "
"%d channels_last upsamples).",
label,
n_norm,
n_up,
)
return vae
python/sglang/multimodal_gen/runtime/platforms/cuda.py
+4 -2
View File
@@ -745,8 +745,8 @@ class CudaPlatformBase(Platform):
@classmethod
def optimize_vae(cls, vae: torch.nn.Module) -> torch.nn.Module:
"""Install the quality-gated FLUX.2 / AutoencoderKL / Wan VAE decoder
fast paths.
"""Install the quality-gated FLUX.2 / AutoencoderKL / Wan / Qwen-Image
VAE decoder fast paths.
Requests with quality="extra-high" or "high" run the fast paths; the
"lossless" default runs the original module path bit-for-bit. See
@@ -758,12 +758,14 @@ class CudaPlatformBase(Platform):
maybe_optimize_flux2_vae,
)
from sglang.multimodal_gen.runtime.models.vaes.wan_vae_cuda_opt import (
maybe_optimize_qwen_image_vae,
maybe_optimize_wan_vae,
)
vae = maybe_optimize_flux2_vae(vae)
vae = maybe_optimize_autoencoder_kl(vae)
vae = maybe_optimize_wan_vae(vae)
vae = maybe_optimize_qwen_image_vae(vae)
except Exception:
logger.warning(
"Failed to apply CUDA VAE optimizations; using the unmodified VAE.",